from flask import Flask, render_template, request, jsonify, send_from_directory
from flask_socketio import SocketIO
from flask_cors import CORS
import mediapipe as mp
import cv2
import numpy as np
import time
import os
import json

app = Flask(__name__)
socketio = SocketIO(app)
CORS(app)

mp_pose = mp.solutions.pose
pose = mp_pose.Pose(static_image_mode=False, min_detection_confidence=0.5, min_tracking_confidence=0.5)

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

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


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

def process_video(video_path):
    cap = cv2.VideoCapture(video_path)
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    
    all_results = []
    previous_keypoints = None
    previous_velocities = None
    previous_time = 0
    
    for frame_number in range(frame_count):
        success, frame = cap.read()
        if not success:
            break
        
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        results = pose.process(frame_rgb)
        
        if results.pose_landmarks:
            landmarks = results.pose_landmarks.landmark
            keypoints = [(lm.x, lm.y, lm.z) for lm in landmarks]
            current_time = frame_number / fps
            delta_time = current_time - previous_time
            
            velocities = []
            accelerations = []
            
            if previous_keypoints:
                for i, lm in enumerate(landmarks):
                    dx = lm.x - previous_keypoints[i][0]
                    dy = lm.y - previous_keypoints[i][1]
                    dz = lm.z - previous_keypoints[i][2]
                    speed = np.sqrt(dx ** 2 + dy ** 2 + dz ** 2) / delta_time if delta_time > 0 else 0
                    velocities.append(speed)
                    acceleration = (speed - previous_velocities[i]) / delta_time if previous_velocities and delta_time > 0 else 0
                    accelerations.append(acceleration)
            
            all_results.append({
                'frame': frame_number,
                'landmarks': keypoints,
                'velocities': velocities,
                'accelerations': accelerations
            })
            
            previous_keypoints = keypoints
            previous_velocities = velocities
            previous_time = current_time
    
    cap.release()
    return all_results, fps

@app.route('/upload_video', methods=['POST'])
def upload_video():
    file = request.files.get('video')
    if file:
        upload_folder = os.path.join(app.root_path, 'static', 'uploads')
        if not os.path.exists(upload_folder):
            os.makedirs(upload_folder)
        video_path = os.path.join(upload_folder, 'temp_video.mp4')
        file.save(video_path)
        
        results, fps = process_video(video_path)
        
        results_path = os.path.join(upload_folder, 'results.json')
        with open(results_path, 'w') as f:
            json.dump({'results': results, 'fps': fps}, f)
        
        return jsonify(success=True, message='Video processed successfully')
    return jsonify(success=False, message='No video file received'), 400

@app.route('/get_results')
def get_results():
    results_path = os.path.join(app.root_path, 'static', 'uploads', 'results.json')
    if os.path.exists(results_path):
        with open(results_path, 'r') as f:
            return jsonify(json.load(f))
    return jsonify(success=False, message='Results not found'), 404

@app.route('/uploads/<path:filename>')
def serve_upload(filename):
    return send_from_directory(os.path.join(app.root_path, 'static', 'uploads'), filename)

if __name__ == '__main__':
    os.makedirs(os.path.join(app.root_path, 'static', 'uploads'), exist_ok=True)
    socketio.run(app, host='0.0.0.0', port=7860, debug=True, allow_unsafe_werkzeug=True)