from flask import Flask, render_template, request, Response
import cv2
import mediapipe as mp
import numpy as np
import os

app = Flask(__name__)

# Initialize Mediapipe Pose
mp_pose = mp.solutions.pose
mp_drawing = mp.solutions.drawing_utils

# Function to calculate the angle between three points
def calculate_angle(a, b, c):
    a = np.array(a)  # First point
    b = np.array(b)  # Middle point
    c = np.array(c)  # Last 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 check shoulder press posture
def is_shoulder_press_correct(landmarks, mp_pose):
    # Get coordinates of shoulder, elbow, and wrist (left arm as example)
    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 angle at the elbow (shoulder, elbow, wrist)
    elbow_angle = calculate_angle(shoulder, elbow, wrist)

    # Check if the motion is vertical (wrist higher than elbow)
    if wrist[1] < elbow[1] and elbow[1] < shoulder[1]:
        # Ensure proper angle range for a shoulder press
        if 160 <= elbow_angle <= 180:
            return "Shoulder Press: Correct", True
        else:
            return "Shoulder Press: Incorrect - Elbow angle", False
    else:
        return "Shoulder Press: Incorrect - Alignment", False

# Video processing function
def process_video(video_path):
    cap = cv2.VideoCapture(video_path)
    count = 0
    rep_started = False

    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:
                break

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

            # Process the image for pose detection
            results = pose.process(image)

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

            # Extract landmarks
            if results.pose_landmarks:
                landmarks = results.pose_landmarks.landmark

                # Check shoulder press posture
                feedback, is_correct = is_shoulder_press_correct(landmarks, mp_pose)

                # Count reps
                if is_correct and not rep_started:
                    rep_started = True
                elif not is_correct and rep_started:
                    rep_started = False
                    count += 1

                # Display feedback
                cv2.putText(image, feedback, (50, 50),
                            cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0) if is_correct else (0, 0, 255), 2, cv2.LINE_AA)

                # Display rep count
                cv2.putText(image, f'Reps: {count}', (50, 100),
                            cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)

                # Draw landmarks
                mp_drawing.draw_landmarks(image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS)

            else:
                # Warn if no landmarks are detected
                cv2.putText(image, "No body detected", (50, 50),
                            cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA)

            # Encode frame for streaming
            ret, buffer = cv2.imencode('.jpg', image)
            frame = buffer.tobytes()
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')

    cap.release()

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/upload', methods=['POST'])
def upload():
    if 'video' not in request.files:
        return "No video file uploaded", 400

    video = request.files['video']
    video_path = os.path.join('uploads', video.filename)
    video.save(video_path)

    return Response(process_video(video_path), mimetype='multipart/x-mixed-replace; boundary=frame')

if __name__ == "__main__":
    app.run(debug=True)