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

# Cache the MediaPipe Pose model
@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

# Function to calculate angle
def calculate_angle(a, b, c):
    a = np.array(a)  # First point
    b = np.array(b)  # Mid point
    c = np.array(c)  # End point
    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

# Function to calculate body angle
def calculate_body_angle(shoulder, hip, ankle):
    shoulder = np.array(shoulder)
    hip = np.array(hip)
    ankle = np.array(ankle)
    radians = np.arctan2(ankle[1] - hip[1], ankle[0] - hip[0]) - np.arctan2(shoulder[1] - hip[1], shoulder[0] - hip[0])
    angle = np.abs(radians * 180.0 / np.pi)
    if angle > 180.0:
        angle = 360 - angle
    return angle

# Function to calculate rep score
def calculate_rep_score(elbow_angles, body_angles):
    ideal_elbow_range = (75, 90)
    ideal_body_range = (0, 10)
    elbow_score = sum(1 for angle in elbow_angles if ideal_elbow_range[0] <= angle <= ideal_elbow_range[1]) / len(elbow_angles) if elbow_angles else 0
    body_score = sum(1 for angle in body_angles if ideal_body_range[0] <= angle <= ideal_body_range[1]) / len(body_angles) if body_angles else 0
    return (elbow_score + body_score) / 2

# Function to generate workout report
def generate_workout_report(rep_scores, form_issues, analysis_time):
    overall_efficiency = sum(rep_scores) / len(rep_scores) if rep_scores else 0
    total_reps = len(rep_scores)
    elbows_not_bending_enough = form_issues['elbows_not_bending_enough']
    body_not_straight = form_issues['body_not_straight']

    report = f"""
    **Workout Report:**
    -----------------
    - **Total Push-ups:** {total_reps}
    - **Overall Workout Efficiency:** {overall_efficiency * 100:.2f}%
    - **Analysis Time:** {analysis_time:.2f} seconds

    **Form Issues:**
    - Elbows not bending enough: {elbows_not_bending_enough} reps ({(elbows_not_bending_enough / total_reps) * 100:.2f}% of reps)
    - Body not straight: {body_not_straight} reps ({(body_not_straight / total_reps) * 100:.2f}% of reps)
    """
    return report

# Load Lottie animation from a JSON file
def load_lottiefile(filepath: str):
    with open(filepath, "r") as f:
        return json.load(f)

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

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

# Path to the demo video
demo_video_path = "W_58.mp4"

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

st.write("Or upload your own video:")
# File uploader for user's 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}")

    # Create placeholders for the three video outputs with titles
    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()

    feedback_placeholder = st.empty()

    rep_scores = []
    current_rep_angles = {'elbow': [], 'body': []}
    form_issues = {
        "elbows_not_bending_enough": 0,
        "body_not_straight": 0
    }
    stage = "UP"
    pushup_count = 0
    start_time = time.time()

    # Initialize form issues for the current rep
    current_rep_issues = {
        "elbows_not_bending_enough": False,
        "body_not_straight": False
    }

    try:
        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:
                try:
                    landmarks = results.pose_landmarks.landmark
                    shoulder = [landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].x * width,
                                landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].y * height]
                    elbow = [landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value].x * width,
                             landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value].y * height]
                    wrist = [landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].x * width,
                             landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].y * height]
                    hip = [landmarks[mp_pose.PoseLandmark.LEFT_HIP.value].x * width,
                           landmarks[mp_pose.PoseLandmark.LEFT_HIP.value].y * height]
                    ankle = [landmarks[mp_pose.PoseLandmark.LEFT_ANKLE.value].x * width,
                             landmarks[mp_pose.PoseLandmark.LEFT_ANKLE.value].y * height]

                    angle_elbow = calculate_angle(shoulder, elbow, wrist)
                    angle_body = calculate_body_angle(shoulder, hip, ankle)

                    current_rep_angles['elbow'].append(angle_elbow)
                    current_rep_angles['body'].append(angle_body)

                    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)
                    mp_drawing.draw_landmarks(guide_image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS)

                    if angle_elbow < 75:
                        current_rep_issues["elbows_not_bending_enough"] = True
                        cv2.line(guide_image, tuple(np.multiply(elbow, [1, 1]).astype(int)),
                                 tuple(np.multiply(wrist, [1, 1]).astype(int)), (0, 255, 255), 5)  # Yellow for elbow issues
                    if angle_body > 10:
                        current_rep_issues["body_not_straight"] = True
                        cv2.line(guide_image, tuple(np.multiply(shoulder, [1, 1]).astype(int)),
                                 tuple(np.multiply(hip, [1, 1]).astype(int)), (0, 0, 255), 5)  # Red for body issues

                    if angle_elbow > 90 and stage != "UP":
                        stage = "UP"
                        # Count issues for the completed rep
                        for issue, occurred in current_rep_issues.items():
                            if occurred:
                                form_issues[issue] += 1
                        # Reset current rep issues
                        current_rep_issues = {k: False for k in current_rep_issues}

                    if angle_elbow < 90 and stage == "UP":
                        stage = "DOWN"
                        pushup_count += 1
                        rep_score = calculate_rep_score(current_rep_angles['elbow'], current_rep_angles['body'])
                        rep_scores.append(rep_score)
                        current_rep_angles = {'elbow': [], 'body': []}

                    # Display the videos
                    original_video.image(image, channels="BGR")
                    points_video.image(points_image, channels="BGR")
                    guide_video.image(guide_image, channels="BGR")
                except AttributeError as e:
                    st.error(f"Error processing frame: {e}")
            else:
                st.warning("No pose landmarks detected in this frame.")
    except Exception as e:
        st.error(f"Error occurred: {e}")
    finally:
        cap.release()

    analysis_time = time.time() - start_time
    report = generate_workout_report(rep_scores, form_issues, analysis_time)
    feedback_placeholder.markdown(report)

    #st.balloons()  # Add a fun animation when the analysis is done