import streamlit as st
import cv2
import mediapipe as mp
import numpy as np
import time
import json

@st.cache_resource
def load_pose_model():
    mp_pose = mp.solutions.pose
    pose = mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5)
    return mp_pose, pose

mp_pose, pose = load_pose_model()
mp_drawing = mp.solutions.drawing_utils

def calculate_angle(a, b, c):
    a = np.array(a)  # First
    b = np.array(b)  # Mid
    c = np.array(c)  # End
    radians = np.arctan2(c[1] - b[1], c[0] - b[0]) - np.arctan2(a[1] - b[1], a[0] - b[0])
    angle = np.abs(radians * 180.0 / np.pi)
    if angle > 180.0:
        angle = 360 - angle
    return angle

def calculate_rep_score(knee_angles, hip_angles, back_angles):
    ideal_knee_range = (90, 110)
    ideal_hip_range = (80, 100)
    ideal_back_range = (70, 90)
    knee_score = sum(1 for angle in knee_angles if ideal_knee_range[0] <= angle <= ideal_knee_range[1]) / len(knee_angles) if knee_angles else 0
    hip_score = sum(1 for angle in hip_angles if ideal_hip_range[0] <= angle <= ideal_hip_range[1]) / len(hip_angles) if hip_angles else 0
    back_score = sum(1 for angle in back_angles if ideal_back_range[0] <= angle <= ideal_back_range[1]) / len(back_angles) if back_angles else 0
    return (knee_score + hip_score + back_score) / 3

def generate_workout_report(rep_scores, form_issues, analysis_time):
    if rep_scores:
        overall_efficiency = sum(rep_scores) / len(rep_scores)
    else:
        overall_efficiency = 0

    total_reps = len(rep_scores)
    
    if total_reps > 0:
        knee_percentage = form_issues['knees_bending_too_much'] / total_reps * 100
        hip_percentage = form_issues['hips_bending_too_much'] / total_reps * 100
        back_percentage = form_issues['back_leaning_too_much'] / total_reps * 100
    else:
        knee_percentage = 0
        hip_percentage = 0
        back_percentage = 0

    report = f"""
    Workout Report:
    ---------------
    Total Squats: {total_reps}
    Overall Workout Efficiency: {overall_efficiency * 100:.2f}%
    Analysis Time: {analysis_time:.2f} seconds
    Form Issues:
    - Knees bending too much: {form_issues['knees_bending_too_much']} reps ({knee_percentage:.2f}% of reps)
    - Hips bending too much: {form_issues['hips_bending_too_much']} reps ({hip_percentage:.2f}% of reps)
    - Back leaning too much: {form_issues['back_leaning_too_much']} reps ({back_percentage:.2f}% of reps)
    """
    return report

def load_lottiefile(filepath: str):
    with open(filepath, "r") as f:
        return json.load(f)

st.markdown("<h1 style='text-align: center;'>Squat Form Analysis</h1>", unsafe_allow_html=True)

col1, col2, col3 = st.columns([1,2,1])
with col2:
    demo_button = st.button("Try Demo")

demo_video_path = "demo.mp4"

video_path = None
if demo_button:
    video_path = demo_video_path
    st.success("Demo video loaded successfully!")

st.write("Or upload your own video:")
uploaded_file = st.file_uploader("Choose a video file", type=["mp4", "mov", "avi"])
if uploaded_file is not None:
    with open("temp_video.mp4", "wb") as f:
        f.write(uploaded_file.getvalue())
    video_path = "temp_video.mp4"
    st.success("Your video uploaded successfully!")

if video_path:
    cap = cv2.VideoCapture(video_path)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    st.write(f"Frames per second: {fps}")

    col1, col2, col3 = st.columns(3)
    with col1:
        st.subheader("Original Video")
        original_video = st.empty()
    with col2:
        st.subheader("Pose Points")
        points_video = st.empty()
    with col3:
        st.subheader("Form Guide")
        guide_video = st.empty()

    squat_count_placeholder = st.empty()
    feedback_placeholder = st.empty()

    rep_scores = []
    current_rep_angles = {'knee': [], 'hip': [], 'back': []}
    form_issues = {
        "knees_bending_too_much": 0,
        "hips_bending_too_much": 0,
        "back_leaning_too_much": 0
    }
    stage = None
    squat_count = 0
    start_time = time.time()

    current_rep_issues = {
        "knees_bending_too_much": False,
        "hips_bending_too_much": False,
        "back_leaning_too_much": False
    }

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        image.flags.writeable = False
        results = pose.process(image)
        image.flags.writeable = True
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

        if results.pose_landmarks is not None:
            landmarks = results.pose_landmarks.landmark
            shoulder = [landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].x * width,
                        landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].y * height]
            hip = [landmarks[mp_pose.PoseLandmark.LEFT_HIP.value].x * width,
                   landmarks[mp_pose.PoseLandmark.LEFT_HIP.value].y * height]
            knee = [landmarks[mp_pose.PoseLandmark.LEFT_KNEE.value].x * width,
                    landmarks[mp_pose.PoseLandmark.LEFT_KNEE.value].y * height]
            ankle = [landmarks[mp_pose.PoseLandmark.LEFT_ANKLE.value].x * width,
                     landmarks[mp_pose.PoseLandmark.LEFT_ANKLE.value].y * height]

            angle_knee = calculate_angle(hip, knee, ankle)
            angle_hip = calculate_angle(shoulder, hip, knee)
            angle_back = calculate_angle(shoulder, hip, ankle)

            current_rep_angles['knee'].append(angle_knee)
            current_rep_angles['hip'].append(angle_hip)
            current_rep_angles['back'].append(angle_back)

            points_image = np.zeros((height, width, 3), dtype=np.uint8)
            guide_image = np.zeros((height, width, 3), dtype=np.uint8)
            mp_drawing.draw_landmarks(points_image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS)

            # Drawing on guide image with smaller markers
            mp_drawing.draw_landmarks(
                guide_image, 
                results.pose_landmarks, 
                mp_pose.POSE_CONNECTIONS,
                mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=1, circle_radius=2),  # Green for upper arm
                mp_drawing.DrawingSpec(color=(255, 0, 0), thickness=1, circle_radius=2)  # Red for lower arm
            )

            if angle_knee < 90:
                current_rep_issues["knees_bending_too_much"] = True
                cv2.line(guide_image, tuple(np.multiply(knee, [1, 1]).astype(int)),
                         tuple(np.multiply(ankle, [1, 1]).astype(int)), (0, 255, 255), 2)
            if angle_hip < 80:
                current_rep_issues["hips_bending_too_much"] = True
                cv2.line(guide_image, tuple(np.multiply(hip, [1, 1]).astype(int)),
                         tuple(np.multiply(knee, [1, 1]).astype(int)), (0, 0, 255), 2)
            if angle_back < 70:
                current_rep_issues["back_leaning_too_much"] = True
                cv2.line(guide_image, tuple(np.multiply(shoulder, [1, 1]).astype(int)),
                         tuple(np.multiply(hip, [1, 1]).astype(int)), (255, 0, 0), 2)

            # Rep counting logic
            if angle_knee > 160 and stage != "UP":
                stage = "UP"
                for issue, occurred in current_rep_issues.items():
                    if occurred:
                        form_issues[issue] += 1
                rep_scores.append(calculate_rep_score(
                    current_rep_angles['knee'],
                    current_rep_angles['hip'],
                    current_rep_angles['back']
                ))
                current_rep_angles = {'knee': [], 'hip': [], 'back': []}
                current_rep_issues = {k: False for k in current_rep_issues}
            elif angle_knee <= 90 and stage == "UP":
                stage = "DOWN"
                squat_count += 1
                squat_count_placeholder.write(f"Squat Count: {squat_count}")

        # Update video streams
        original_video.image(frame, channels="BGR", use_column_width=True)
        points_video.image(points_image, channels="BGR", use_column_width=True)
        guide_video.image(guide_image, channels="BGR", use_column_width=True)

    cap.release()
    end_time = time.time()
    analysis_time = end_time - start_time

    # Generate and display the report
    report = generate_workout_report(rep_scores, form_issues, analysis_time)
    feedback_placeholder.text(report)