import cv2
import mediapipe as mp
import numpy as np
import gradio as gr

mp_pose = mp.solutions.pose
mp_drawing = mp.solutions.drawing_utils

# Custom angle calculation for bench press (ensuring angle is within 0-180)
def calculate_angle_bp(a, b, c):
    a, b, c = np.array(a), np.array(b), 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))
    return angle if angle <= 180 else 360 - angle

def check_bench_press_feedback(landmarks):
    target_angle = 90
    margin = 10
    left_shoulder = [landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].x,
                     landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].y]
    left_elbow = [landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value].x,
                  landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value].y]
    left_wrist = [landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].x,
                  landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].y]
    right_shoulder = [landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER.value].x,
                      landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER.value].y]
    right_elbow = [landmarks[mp_pose.PoseLandmark.RIGHT_ELBOW.value].x,
                   landmarks[mp_pose.PoseLandmark.RIGHT_ELBOW.value].y]
    right_wrist = [landmarks[mp_pose.PoseLandmark.RIGHT_WRIST.value].x,
                   landmarks[mp_pose.PoseLandmark.RIGHT_WRIST.value].y]
    
    left_angle = calculate_angle_bp(left_shoulder, left_elbow, left_wrist)
    right_angle = calculate_angle_bp(right_shoulder, right_elbow, right_wrist)
    avg_angle = (left_angle + right_angle) / 2.0
    diff = abs(avg_angle - target_angle)
    accuracy = max(0, min(100, int((1 - diff / margin) * 100))) if diff <= margin else 0
    
    if (target_angle - margin) <= left_angle <= (target_angle + margin) and \
       (target_angle - margin) <= right_angle <= (target_angle + margin):
        feedback = "Correct Bench Press Form"
    else:
        feedback = "Incorrect Bench Press Form: "
        if left_angle < target_angle - margin:
            feedback += "Left elbow too bent. "
        elif left_angle > target_angle + margin:
            feedback += "Left elbow not bent enough. "
        if right_angle < target_angle - margin:
            feedback += "Right elbow too bent. "
        elif right_angle > target_angle + margin:
            feedback += "Right elbow not bent enough. "
    return feedback, accuracy

def draw_accuracy_bar(image, accuracy):
    bar_x, bar_y = 50, image.shape[0] - 50
    bar_width, bar_height = 200, 20
    fill_width = int((accuracy / 100) * bar_width)
    color = (0, 255, 0) if accuracy >= 80 else (0, 0, 255) if accuracy < 50 else (0, 255, 255)
    cv2.rectangle(image, (bar_x, bar_y), (bar_x+bar_width, bar_y+bar_height), (200,200,200), 2)
    cv2.rectangle(image, (bar_x, bar_y), (bar_x+fill_width, bar_y+bar_height), color, -1)
    cv2.putText(image, f"Accuracy: {accuracy}%", (bar_x, bar_y-10),
                cv2.FONT_HERSHEY_DUPLEX, 0.6, (255,255,255), 2)

def analyze_bench_press(video_path):
    cap = cv2.VideoCapture(video_path)
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS) if cap.get(cv2.CAP_PROP_FPS) > 0 else 30

    output_video = "output_bench_press.mp4"
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_video, fourcc, fps, (frame_width, frame_height))
    
    # Use a fresh pose instance in a context for bench press analysis
    with mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5) as pose_instance:
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            results = pose_instance.process(image_rgb)
            annotated_frame = frame.copy()
            if results.pose_landmarks:
                mp_drawing.draw_landmarks(annotated_frame, results.pose_landmarks, mp_pose.POSE_CONNECTIONS)
                landmarks = results.pose_landmarks.landmark
                feedback, accuracy = check_bench_press_feedback(landmarks)
                draw_accuracy_bar(annotated_frame, accuracy)
                text_color = (0, 255, 0) if "Correct" in feedback else (0, 0, 255)
                cv2.putText(annotated_frame, feedback, (50, 50),
                            cv2.FONT_HERSHEY_COMPLEX, 1, text_color, 3)
            out.write(annotated_frame)
    cap.release()
    out.release()
    return output_video

gr.Interface(
    fn=analyze_bench_press,
    inputs=gr.Video(),
    outputs=gr.Video(),
    title="Bench Press Form Analyzer",
    description="Upload a video of your bench press and receive form feedback!"
).launch()