# Streamlit App for Workout Tracker
import streamlit as st
import cv2
import mediapipe as mp
import numpy as np
import time
from sklearn.ensemble import IsolationForest

# Title and Introduction
st.title("Muscle Memory")
st.markdown("""
Welcome to the **Workout Tracker App**! 
Select your desired workout below, and the app will guide you through the exercise with real-time feedback.
""")

# Workout Options
st.header("Choose Your Workout")
workout_option = st.selectbox(
    "Available Workouts:",
    ["Bicep Curl", "Lateral Raise", "Shoulder Press"]
)

# Mediapipe utilities
mp_drawing = mp.solutions.drawing_utils
mp_pose = mp.solutions.pose


# Helper Functions
def calculate_angle(a, b, c):
    """Calculate the angle between three points."""
    a = np.array(a)
    b = np.array(b)
    c = np.array(c)

    radians = np.arctan2(c[1] - b[1], c[0] - b[0]) - np.arctan2(a[1] - b[1], a[0] - b[0])
    angle = np.abs(np.degrees(radians))
    if angle > 180.0:
        angle = 360 - angle
    return angle


def draw_text_with_background(image, text, position, font, font_scale, color, thickness, bg_color, padding=10):
    """Draw text with a background on an image."""
    text_size = cv2.getTextSize(text, font, font_scale, thickness)[0]
    text_x, text_y = position
    box_coords = (
        (text_x - padding, text_y - padding),
        (text_x + text_size[0] + padding, text_y + text_size[1] + padding),
    )
    cv2.rectangle(image, box_coords[0], box_coords[1], bg_color, cv2.FILLED)
    cv2.putText(image, text, (text_x, text_y + text_size[1]), font, font_scale, color, thickness)


# Main Function to Track Workout
def track_workout():
    cap = cv2.VideoCapture(0)
    counter = 0  # Rep counter
    stage = None  # Movement stage
    max_reps = 10  # Max reps
    feedback = ""  # Real-time feedback for the video feed

    with mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5) as pose:
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                st.error("Failed to access the webcam. Please ensure your webcam is enabled.")
                break

            # Convert frame to RGB
            image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            image.flags.writeable = False
            results = pose.process(image)

            # Convert back to BGR
            image.flags.writeable = True
            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

            # Check if pose landmarks are detected
            if results.pose_landmarks:
                landmarks = results.pose_landmarks.landmark

                # Extract key joints
                shoulder = [
                    landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].x,
                    landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].y,
                ]
                elbow = [
                    landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value].x,
                    landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value].y,
                ]
                wrist = [
                    landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].x,
                    landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].y,
                ]

                # Calculate the angle
                angle = calculate_angle(shoulder, elbow, wrist)

                # Stage logic for counting reps
                if angle > 160 and stage != "down":
                    stage = "down"
                elif angle < 40 and stage == "down":
                    stage = "up"
                    counter += 1

                    # Feedback
                    if counter == max_reps:
                        feedback = "Workout Complete!"
                        break
                    else:
                        feedback = f"Rep {counter} completed!"

                # Draw the feedback on the frame
                draw_text_with_background(image, f"Reps: {counter}", (50, 50),
                                          cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, (0, 0, 0))
                draw_text_with_background(image, feedback, (50, 100),
                                          cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, (0, 0, 0))

                # Draw landmarks
                mp_drawing.draw_landmarks(
                    image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS,
                    mp_drawing.DrawingSpec(color=(245, 117, 66), thickness=2, circle_radius=2),
                    mp_drawing.DrawingSpec(color=(245, 66, 230), thickness=2, circle_radius=2)
                )

            # Display the video feed
            cv2.imshow("Workout Tracker", image)

            # Break if 'q' is pressed
            if cv2.waitKey(10) & 0xFF == ord("q"):
                break

        cap.release()
        cv2.destroyAllWindows()


# Button to Start the Workout
if st.button("Start Workout"):
    st.write(f"Starting {workout_option} workout...")
    track_workout()

# Footer
st.markdown("""
---
**Note**: Press "q" in the webcam feed to stop the workout. Ensure your webcam is enabled.
""")