exercise_tracker.py

import time
from enum import Enum
import math
import csv
import os
from datetime import datetime
import numpy as np
from pose_utils import calculate_angle, get_landmark_coordinates

class ExerciseType(Enum):
    BICEP_CURL = 1
    SQUAT = 2
    PUSH_UP = 3
    LUNGE = 4

class ExerciseTracker:
    def __init__(self):
        self.exercise_type = ExerciseType.BICEP_CURL
        self.previous_exercise_type = None
        
        # Store exercise data for each type
        self.exercise_data = {
            ExerciseType.BICEP_CURL: {
                "left_counter": 0,
                "right_counter": 0,
                "set_counter": 1,
                "total_reps": 0,
                "stage": {"left": None, "right": None},
                "rep_start_time": time.time(),
                "rep_speeds": [],
                "rep_counted": {"left": False, "right": False}  # Track if rep was already counted
            },
            ExerciseType.SQUAT: {
                "counter": 0,
                "set_counter": 1,
                "total_reps": 0,
                "stage": None,
                "rep_start_time": time.time(),
                "rep_speeds": [],
                "rep_counted": False  # Track if rep was already counted
            },
            ExerciseType.PUSH_UP: {
                "counter": 0,
                "set_counter": 1,
                "total_reps": 0,
                "stage": None,
                "rep_start_time": time.time(),
                "rep_speeds": [],
                "rep_counted": False  # Track if rep was already counted
            },
            ExerciseType.LUNGE: {
                "counter": 0,
                "set_counter": 1,
                "total_reps": 0,
                "stage": None,
                "rep_start_time": time.time(),
                "rep_speeds": [],
                "rep_counted": False  # Track if rep was already counted
            }
        }
        
        self.start_time = time.time()
        self.feedback = ""
        self.previous_feedback = ""  # Track previous feedback to avoid repetition
        self.feedback_color = (0, 0, 0)  # Default black
        self.is_resting = False
        self.rest_start_time = 0
        self.rest_duration = 60  # 60 seconds rest between sets
        self.rep_progress = 0  # For progress bar
        self.symmetry_score = 100  # For circular gauge
        self.last_spoken_time = 0  # Track when feedback was last spoken
        self.feedback_cooldown = 3  # Minimum seconds between feedback messages
        
    def get_current_data(self):
        """Get the data for the current exercise type"""
        return self.exercise_data[self.exercise_type]

    def update_bicep_curl(self, landmarks, image_shape):
        current_data = self.get_current_data()
        # Determine left/right based on screen position
        screen_left, screen_right = self.determine_left_right(landmarks, image_shape)
        # Track both arms using screen-based left/right
        left_metrics = self._process_arm(landmarks, screen_left, image_shape, current_data)
        right_metrics = self._process_arm(landmarks, screen_right, image_shape, current_data)
        # Calculate symmetry (difference between arms) with safe bounds
        angle_diff = abs(left_metrics["angle"] - right_metrics["angle"])
        self.symmetry_score = max(0, 100 - min(100, angle_diff * 2))
        # Calculate rep progress for progress bar (normalized 0-1)
        if left_metrics["angle"] < 30 or right_metrics["angle"] < 30:
            # Normalize progress based on the arm that's closer to completion
            progress_left = max(0, 1 - left_metrics["angle"] / 30)
            progress_right = max(0, 1 - right_metrics["angle"] / 30)
            self.rep_progress = max(progress_left, progress_right)
        else:
            self.rep_progress = 0
        # Calculate speed (reps per minute)
        current_time = time.time()
        if left_metrics["completed_rep"] or right_metrics["completed_rep"]:
            rep_time = current_time - current_data["rep_start_time"]
            current_data["rep_speeds"].append(rep_time)
            current_data["rep_start_time"] = current_time
            # Keep only last 5 reps for average
            if len(current_data["rep_speeds"]) > 5:
                current_data["rep_speeds"].pop(0)
        # Form feedback with color coding
        self.feedback = ""
        self.feedback_color = (0, 255, 0)  # Default to green (good form)
        if left_metrics["angle"] > 100 and current_data["stage"][screen_left] == "down":
            self.feedback = "Left arm: Bend your elbow more for full range"
            self.feedback_color = (0, 0, 255)  # Red for bad form
        elif right_metrics["angle"] > 100 and current_data["stage"][screen_right] == "down":
            self.feedback = "Right arm: Bend your elbow more for full range"
            self.feedback_color = (0, 0, 255)  # Red for bad form
        elif angle_diff > 15:
            self.feedback = "Try to keep both arms symmetrical"
            self.feedback_color = (0, 165, 255)  # Orange for warning
        elif left_metrics["angle"] < 30 and right_metrics["angle"] < 30:
            self.feedback = "Good form! Full range of motion"
            self.feedback_color = (0, 255, 0)  # Green for good form
        return {
            "left_angle": left_metrics["angle"],
            "right_angle": right_metrics["angle"],
            "symmetry": self.symmetry_score,
            "left_elbow_pos": left_metrics["elbow_pos"],
            "right_elbow_pos": right_metrics["elbow_pos"]
        }

    def _process_arm(self, landmarks, side="left", image_shape=None, current_data=None):
        shoulder = get_landmark_coordinates(landmarks, "SHOULDER", side)
        elbow = get_landmark_coordinates(landmarks, "ELBOW", side)
        wrist = get_landmark_coordinates(landmarks, "WRIST", side)
        if shoulder and elbow and wrist:
            angle = calculate_angle(shoulder[:2], elbow[:2], wrist[:2])
            # Use screen-based side for stage tracking
            stage_key = "left" if side == "left" else "right"
            stage = current_data["stage"][stage_key]
            counter = current_data["left_counter"] if stage_key == "left" else current_data["right_counter"]
            completed_rep = False
            # Curl counter logic with state tracking to prevent multiple counts
            if angle > 160:
                current_data["stage"][stage_key] = "down"
                current_data["rep_counted"][stage_key] = False  # Reset rep counted flag when arm is extended
            if angle < 30 and current_data["stage"][stage_key] == 'down' and not current_data["rep_counted"][stage_key]:
                current_data["stage"][stage_key] = "up"
                counter += 1
                current_data["total_reps"] += 1
                completed_rep = True
                current_data["rep_counted"][stage_key] = True  # Mark rep as counted
                print(f"{stage_key.capitalize()} counter: {counter}")
            # Update the correct counter
            if stage_key == "left":
                current_data["left_counter"] = counter
            else:
                current_data["right_counter"] = counter
            # Calculate elbow position with dynamic image shape
            h, w = image_shape[:2] if image_shape is not None else (640, 480)
            elbow_pos = tuple(np.multiply(elbow[:2], [w, h]).astype(int))
            return {
                "angle": angle,
                "completed_rep": completed_rep,
                "elbow_pos": elbow_pos
            }
        return {"angle": 0, "completed_rep": False, "elbow_pos": (0, 0)}
    
    def update_squat(self, landmarks, image_shape):
        current_data = self.get_current_data()
        
        # Get relevant landmarks
        hip_left = get_landmark_coordinates(landmarks, "HIP", "left")
        knee_left = get_landmark_coordinates(landmarks, "KNEE", "left")
        ankle_left = get_landmark_coordinates(landmarks, "ANKLE", "left")
        
        hip_right = get_landmark_coordinates(landmarks, "HIP", "right")
        knee_right = get_landmark_coordinates(landmarks, "KNEE", "right")
        ankle_right = get_landmark_coordinates(landmarks, "ANKLE", "right")
        
        if hip_left and knee_left and ankle_left and hip_right and knee_right and ankle_right:
            # Calculate angles for both legs
            left_angle = calculate_angle(hip_left[:2], knee_left[:2], ankle_left[:2])
            right_angle = calculate_angle(hip_right[:2], knee_right[:2], ankle_right[:2])
            avg_angle = (left_angle + right_angle) / 2
            
            # Calculate rep progress for progress bar (normalized 0-1)
            if avg_angle < 160:
                self.rep_progress = 1 - (avg_angle / 160)
            else:
                self.rep_progress = 0
            
            # Calculate symmetry with safe bounds
            angle_diff = abs(left_angle - right_angle)
            self.symmetry_score = max(0, 100 - min(100, angle_diff))
            
            # Squat counter logic with state tracking to prevent multiple counts
            if avg_angle > 160:
                current_data["stage"] = "up"
                current_data["rep_counted"] = False  # Reset rep counted flag when standing
            if avg_angle < 90 and current_data["stage"] == "up" and not current_data["rep_counted"]:
                current_data["stage"] = "down"
                current_data["counter"] += 1
                current_data["total_reps"] += 1
                current_data["rep_counted"] = True  # Mark rep as counted
                print(f"Squat counter: {current_data['counter']}")
            
            # Form feedback with color coding
            self.feedback = ""
            self.feedback_color = (0, 255, 0)  # Default to green
            
            if angle_diff > 15:
                self.feedback = "Keep your knees aligned for better form"
                self.feedback_color = (0, 0, 255)  # Red for bad form
            elif avg_angle < 90 and current_data["stage"] == "down":
                self.feedback = "Good depth! Now push up"
                self.feedback_color = (0, 255, 0)  # Green for good form
            elif avg_angle > 100 and current_data["stage"] == "up":
                self.feedback = "Go deeper for full range of motion"
                self.feedback_color = (0, 165, 255)  # Orange for warning
            
            # Calculate knee position with dynamic image shape
            h, w = image_shape[:2] if image_shape is not None else (640, 480)
            knee_pos = tuple(np.multiply(knee_left[:2], [w, h]).astype(int))
            
            return {
                "left_angle": left_angle,
                "right_angle": right_angle,
                "knee_pos": knee_pos
            }
        
        return {"left_angle": 0, "right_angle": 0, "knee_pos": (0, 0)}
    
    def update_push_up(self, landmarks, image_shape):
        current_data = self.get_current_data()
        
        # Get relevant landmarks
        shoulder_left = get_landmark_coordinates(landmarks, "SHOULDER", "left")
        elbow_left = get_landmark_coordinates(landmarks, "ELBOW", "left")
        wrist_left = get_landmark_coordinates(landmarks, "WRIST", "left")
        
        shoulder_right = get_landmark_coordinates(landmarks, "SHOULDER", "right")
        elbow_right = get_landmark_coordinates(landmarks, "ELBOW", "right")
        wrist_right = get_landmark_coordinates(landmarks, "WRIST", "right")
        
        if shoulder_left and elbow_left and wrist_left and shoulder_right and elbow_right and wrist_right:
            # Calculate angles for both arms
            left_angle = calculate_angle(shoulder_left[:2], elbow_left[:2], wrist_left[:2])
            right_angle = calculate_angle(shoulder_right[:2], elbow_right[:2], wrist_right[:2])
            avg_angle = (left_angle + right_angle) / 2
            
            # Calculate rep progress for progress bar (normalized 0-1)
            if avg_angle < 160:
                self.rep_progress = 1 - (avg_angle / 160)
            else:
                self.rep_progress = 0
            
            # Calculate symmetry with safe bounds
            angle_diff = abs(left_angle - right_angle)
            self.symmetry_score = max(0, 100 - min(100, angle_diff))
            
            # Push-up counter logic with state tracking to prevent multiple counts
            if avg_angle > 160:
                current_data["stage"] = "up"
                current_data["rep_counted"] = False  # Reset rep counted flag when arms are extended
            if avg_angle < 90 and current_data["stage"] == "up" and not current_data["rep_counted"]:
                current_data["stage"] = "down"
                current_data["counter"] += 1
                current_data["total_reps"] += 1
                current_data["rep_counted"] = True  # Mark rep as counted
                print(f"Push-up counter: {current_data['counter']}")
            
            # Form feedback with color coding
            self.feedback = ""
            self.feedback_color = (0, 255, 0)  # Default to green
            
            back_straight = self._check_back_straight(landmarks)
            if not back_straight:
                self.feedback = "Keep your back straight for better form"
                self.feedback_color = (0, 0, 255)  # Red for bad form
            elif avg_angle < 90 and current_data["stage"] == "down":
                self.feedback = "Good depth! Now push up"
                self.feedback_color = (0, 255, 0)  # Green for good form
            elif avg_angle > 100 and current_data["stage"] == "up":
                self.feedback = "Go deeper for full range of motion"
                self.feedback_color = (0, 165, 255)  # Orange for warning
            
            # Calculate elbow position with dynamic image shape
            h, w = image_shape[:2] if image_shape is not None else (640, 480)
            elbow_pos = tuple(np.multiply(elbow_left[:2], [w, h]).astype(int))
            
            return {
                "left_angle": left_angle,
                "right_angle": right_angle,
                "elbow_pos": elbow_pos
            }
        
        return {"left_angle": 0, "right_angle": 0, "elbow_pos": (0, 0)}
    
    def update_lunge(self, landmarks, image_shape):
        current_data = self.get_current_data()
        
        # Get relevant landmarks for both legs
        hip_left = get_landmark_coordinates(landmarks, "HIP", "left")
        knee_left = get_landmark_coordinates(landmarks, "KNEE", "left")
        ankle_left = get_landmark_coordinates(landmarks, "ANKLE", "left")
        
        hip_right = get_landmark_coordinates(landmarks, "HIP", "right")
        knee_right = get_landmark_coordinates(landmarks, "KNEE", "right")
        ankle_right = get_landmark_coordinates(landmarks, "ANKLE", "right")
        
        # Get shoulder landmarks for hip angle calculation
        shoulder_left = get_landmark_coordinates(landmarks, "SHOULDER", "left")
        shoulder_right = get_landmark_coordinates(landmarks, "SHOULDER", "right")
        
        if (hip_left and knee_left and ankle_left and hip_right and knee_right and ankle_right and
            shoulder_left and shoulder_right):
            
            # Calculate angles for both legs
            left_knee_angle = calculate_angle(hip_left[:2], knee_left[:2], ankle_left[:2])
            right_knee_angle = calculate_angle(hip_right[:2], knee_right[:2], ankle_right[:2])
            
            # Calculate hip angles to determine which leg is forward
            left_hip_angle = calculate_angle(knee_left[:2], hip_left[:2], shoulder_left[:2])
            right_hip_angle = calculate_angle(knee_right[:2], hip_right[:2], shoulder_right[:2])
            
            # Determine which leg is forward (the one with the smaller hip angle)
            if left_hip_angle < right_hip_angle:
                # Left leg is forward
                forward_angle = left_knee_angle
                back_angle = right_knee_angle
                h, w = image_shape[:2] if image_shape is not None else (640, 480)
                forward_knee_pos = tuple(np.multiply(knee_left[:2], [w, h]).astype(int))
            else:
                # Right leg is forward
                forward_angle = right_knee_angle
                back_angle = left_knee_angle
                h, w = image_shape[:2] if image_shape is not None else (640, 480)
                forward_knee_pos = tuple(np.multiply(knee_right[:2], [w, h]).astype(int))
            
            # Calculate rep progress for progress bar (normalized 0-1)
            if forward_angle < 160:
                self.rep_progress = 1 - (forward_angle / 160)
            else:
                self.rep_progress = 0
            
            # Calculate symmetry with safe bounds
            angle_diff = abs(left_knee_angle - right_knee_angle)
            self.symmetry_score = max(0, 100 - min(100, angle_diff))
            
            # Lunge counter logic with state tracking to prevent multiple counts
            if forward_angle > 160:
                current_data["stage"] = "up"
                current_data["rep_counted"] = False  # Reset rep counted flag when standing
            if forward_angle < 90 and current_data["stage"] == "up" and back_angle > 160 and not current_data["rep_counted"]:
                current_data["stage"] = "down"
                current_data["counter"] += 1
                current_data["total_reps"] += 1
                current_data["rep_counted"] = True  # Mark rep as counted
                print(f"Lunge counter: {current_data['counter']}")
            
            # Form feedback with color coding
            self.feedback = ""
            self.feedback_color = (0, 255, 0)  # Default to green
            
            if forward_angle < 90 and back_angle < 160:
                self.feedback = "Straighten your back leg more"
                self.feedback_color = (0, 0, 255)  # Red for bad form
            elif forward_angle < 90 and back_angle > 160:
                self.feedback = "Good form! Now push back up"
                self.feedback_color = (0, 255, 0)  # Green for good form
            elif forward_angle > 100 and current_data["stage"] == "up":
                self.feedback = "Step forward into a deeper lunge"
                self.feedback_color = (0, 165, 255)  # Orange for warning
            elif angle_diff > 20:
                self.feedback = "Keep your knees aligned with your ankles"
                self.feedback_color = (0, 0, 255)  # Red for bad form
            
            return {
                "left_angle": left_knee_angle,
                "right_angle": right_knee_angle,
                "knee_pos": forward_knee_pos
            }
        
        return {"left_angle": 0, "right_angle": 0, "knee_pos": (0, 0)}
    
    def _check_back_straight(self, landmarks):
        """Check for back straightness using both sides for better accuracy"""
        # Check both sides and average for better accuracy
        left_straight = self._check_side_straight(landmarks, "left")
        right_straight = self._check_side_straight(landmarks, "right")
        
        # If one side is not detected, use the other
        if left_straight is None:
            return right_straight
        if right_straight is None:
            return left_straight
            
        # Return True only if both sides are straight
        return left_straight and right_straight
    
    def _check_side_straight(self, landmarks, side="left"):
        """Check if back is straight for one side"""
        shoulder = get_landmark_coordinates(landmarks, "SHOULDER", side)
        hip = get_landmark_coordinates(landmarks, "HIP", side)
        ankle = get_landmark_coordinates(landmarks, "ANKLE", side)
        
        if shoulder and hip and ankle:
            # Calculate angle between shoulder, hip, and ankle
            angle = calculate_angle(shoulder[:2], hip[:2], ankle[:2])
            # If angle is close to 180 degrees, back is straight
            return abs(angle - 180) < 15
        
        return None
    
    def check_set_completion(self):
        current_data = self.get_current_data()
        # Check if we've completed a set (10 reps) and start rest timer
        if current_data["total_reps"] >= 10 and not self.is_resting:
            self.is_resting = True
            self.rest_start_time = time.time()
            current_data["set_counter"] += 1
            current_data["total_reps"] = 0
            return True
        return False
    
    def get_rest_time_remaining(self):
        if self.is_resting:
            elapsed = time.time() - self.rest_start_time
            remaining = max(0, self.rest_duration - elapsed)
            if remaining <= 0:
                self.is_resting = False
            return remaining
        return 0
    
    def change_exercise(self, new_exercise):
        """Change to a new exercise type"""
        if new_exercise != self.exercise_type:
            self.previous_exercise_type = self.exercise_type
            self.exercise_type = new_exercise
            print(f"Switched to {new_exercise.name.replace('_', ' ').title()}")
    
    def save_to_csv(self, filename="workout_data.csv"):
        """Save all exercise data to a CSV file"""
        try:
            file_exists = os.path.isfile(filename)
            
            with open(filename, mode='a', newline='') as file:
                writer = csv.writer(file)
                
                # Write header if file doesn't exist
                if not file_exists:
                    writer.writerow(["Timestamp", "Exercise", "Left Reps", "Right Reps", "Total Reps", "Sets", "Workout Duration"])
                
                # Write data for each exercise
                timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")  # Fixed date format
                workout_duration = time.time() - self.start_time
                
                for exercise_type, data in self.exercise_data.items():
                    exercise_name = exercise_type.name.replace('_', ' ').title()
                    
                    if exercise_type == ExerciseType.BICEP_CURL:
                        writer.writerow([
                            timestamp,
                            exercise_name,
                            data["left_counter"],
                            data["right_counter"],
                            data["total_reps"],
                            data["set_counter"],
                            f"{int(workout_duration // 60)}:{int(workout_duration % 60):02d}"
                        ])
                    else:
                        writer.writerow([
                            timestamp,
                            exercise_name,
                            "-",  # Left reps not applicable
                            "-",  # Right reps not applicable
                            data["total_reps"],
                            data["set_counter"],
                            f"{int(workout_duration // 60)}:{int(workout_duration % 60):02d}"
                        ])

                # Fixed divider row
                writer.writerow(["---" for _ in range(7)])  # Simple divider with 7 columns
                
            print(f"Workout data saved to {filename}")
            return True
        except Exception as e:
            print(f"Error saving to CSV: {e}")
            return False
    
    def get_tutorial_text(self):
        """Get tutorial text for the current exercise"""
        if self.exercise_type == ExerciseType.BICEP_CURL:
            return "For bicep curls, stand straight with your arms fully extended. Hold weights in each hand with palms facing forward. Curl the weights while keeping your elbows close to your body. Squeeze your biceps at the top, then slowly lower back to starting position."
        elif self.exercise_type == ExerciseType.SQUAT:
            return "For squats, stand with feet shoulder-width apart. Lower your body by bending your knees and pushing your hips back. Keep your chest up and back straight. Go as low as you can while keeping good form, then push through your heels to return to standing."
        elif self.exercise_type == ExerciseType.PUSH_UP:
            return "For push-ups, start in a plank position with hands slightly wider than shoulder-width. Lower your body until your chest nearly touches the floor, keeping your elbows close to your body. Push back up to the starting position while maintaining a straight line from head to heels."
        elif self.exercise_type == ExerciseType.LUNGE:
            return "For lunges, stand with feet hip-width apart. Step forward with one leg, lowering your hips until both knees are bent at a 90-degree angle. Your front knee should be directly above your ankle, and your back knee should hover just above the ground. Push back to the starting position."
        return "No tutorial available for this exercise."
    
    def get_form_feedback(self):
        """Get form feedback with suggestions for improvement"""
        feedback_text = self.feedback
        
        if self.exercise_type == ExerciseType.BICEP_CURL:
            if "Bend your elbow" in self.feedback:
                feedback_text += ". Keep your upper arm stationary and only move your forearm."
            elif "symmetrical" in self.feedback:
                feedback_text += ". Focus on moving both arms at the same pace."
        
        elif self.exercise_type == ExerciseType.SQUAT:
            if "knees aligned" in self.feedback:
                feedback_text += ". Keep your knees in line with your feet."
            elif "deeper" in self.feedback:
                feedback_text += ". Try to get your thighs parallel to the ground."
        
        elif self.exercise_type == ExerciseType.PUSH_UP:
            if "back straight" in self.feedback:
                feedback_text += ". Engage your core and keep your body in a straight line."
            elif "deeper" in self.feedback:
                feedback_text += ". Lower your body until your chest nearly touches the floor."
        
        elif self.exercise_type == ExerciseType.LUNGE:
            if "back leg" in self.feedback:
                feedback_text += ". Keep your back leg straight but not locked."
            elif "knees aligned" in self.feedback:
                feedback_text += ". Your front knee should be directly above your ankle."
        
        return feedback_text
    
    def should_speak_feedback(self):
        """Check if feedback should be spoken (avoid repetition)"""
        current_time = time.time()
        if (self.feedback and self.feedback != self.previous_feedback and 
            current_time - self.last_spoken_time > self.feedback_cooldown):
            self.previous_feedback = self.feedback
            self.last_spoken_time = current_time
            return True
        return False
    
    def determine_left_right(self, landmarks, image_shape):
        h, w = image_shape[:2] if image_shape is not None else (640, 480)
        
        # Get shoulder positions
        left_shoulder = get_landmark_coordinates(landmarks, "SHOULDER", "left")
        right_shoulder = get_landmark_coordinates(landmarks, "SHOULDER", "right")
        
        if left_shoulder and right_shoulder:
            # Convert to pixel coordinates
            left_x = int(left_shoulder[0] * w)
            right_x = int(right_shoulder[0] * w)
            
            # The shoulder with smaller x-coordinate is left from camera's perspective
            if left_x < right_x:
                return "left", "right"  # Correct orientation
            else:
                return "right", "left"  # Flipped orientation
        return "left", "right"  # Default assumption