from flask import Flask, render_template, Response, request, jsonify, send_file
from flask_socketio import SocketIO
import cv2
import json
import mediapipe as mp
import numpy as np
from scipy.spatial import procrustes
import base64
import os

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

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

image_name = "Camel"
image = None

# Extract landmarks from JSON data
landmarks_from_json = []
the_landmarks = None
dataset = {"name": "", "ket": ""}
similarity = 0
all_data = []

def load_image_and_landmarks(image_name):
    global image, landmarks_from_json, the_landmarks, all_data
    landmarks_from_json = []  # Clear previous landmarks
    # Load JSON data
    with open('data_yoga.json') as f:
        data = json.load(f)
        all_data = data

    # Load the image and landmarks
    for the_data in data:
        if the_data['name'] == image_name:
            for lm in the_data['landmarks']:
                landmarks_from_json.append([lm['coordinates'][0], lm['coordinates'][1]])
            the_landmarks = the_data['landmarks']
            image = cv2.imread(the_data['image_name'])
            dataset["name"] = the_data['name']
            dataset["ket"] = the_data['ket']

# Define the function to calculate the color based on similarity
def calculate_color(similarity):
    if similarity < 70:
        return (0, 0, 255)  # Red
    else:
        normalized_similarity = (similarity - 55) / 45  # Normalize between 0 and 1 for values 71 to 100
        red = int((1 - normalized_similarity) * 255)
        green = int(normalized_similarity * 255)
        return (0, green, red)

def resample_landmarks(landmarks, target_length):
    idxs = np.linspace(0, len(landmarks) - 1, target_length).astype(int)
    return [landmarks[i] for i in idxs]

def calculate_similarity(landmarks1, landmarks2):
    if not landmarks1 or not landmarks2:
        return 0
    
    len1 = len(landmarks1)
    len2 = len(landmarks2)
    
    # Resample landmarks to ensure they have the same number of points
    if len1 != len2:
        if len1 < len2:
            landmarks2 = resample_landmarks(landmarks2, len1)
        else:
            landmarks1 = resample_landmarks(landmarks1, len2)
    
    # Convert to numpy arrays
    landmarks1 = np.array(landmarks1)
    landmarks2 = np.array(landmarks2)
    
    # Normalize landmarks by removing the mean
    norm_landmarks1 = landmarks1 - np.mean(landmarks1, axis=0)
    norm_landmarks2 = landmarks2 - np.mean(landmarks2, axis=0)
    
    # Perform Procrustes analysis to align the shapes
    _, mtx1, mtx2 = procrustes(norm_landmarks1, norm_landmarks2)
    
    # Calculate the Euclidean distances between corresponding landmarks
    dists = np.linalg.norm(mtx1 - mtx2, axis=1)
    
    # Calculate the similarity as the inverse of the average distance
    avg_dist = np.mean(dists)
    similarity = max(0, 1 - avg_dist)
    
    # Scale the similarity to a percentage
    similarity_percentage = similarity * 100
    
    return similarity_percentage

def draw_landmarks(image, landmarks):
    global similarity
    annotated_image = image.copy()
    for landmark in landmarks:
        landmark_x = int(landmark['coordinates'][0] * annotated_image.shape[1])
        landmark_y = int(landmark['coordinates'][1] * annotated_image.shape[0])
        cv2.circle(annotated_image, (landmark_x, landmark_y), 5, (0, 255, 0), -1)

    connections = mp_pose.POSE_CONNECTIONS
    for connection in connections:
        start_idx = connection[0]
        end_idx = connection[1]
        if 0 <= start_idx < len(landmarks) and 0 <= end_idx < len(landmarks):
            start_landmark = landmarks[start_idx]['coordinates']
            end_landmark = landmarks[end_idx]['coordinates']
            start_x = int(start_landmark[0] * annotated_image.shape[1])
            start_y = int(start_landmark[1] * annotated_image.shape[0])
            end_x = int(end_landmark[0] * annotated_image.shape[1])
            end_y = int(end_landmark[1] * annotated_image.shape[0])
            cv2.line(annotated_image, (start_x, start_y), (end_x, end_y), (0, 255, 0), 2)

    return annotated_image

def generate_frames():
    global similarity
    cap = cv2.VideoCapture(0)
    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

            image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            image_rgb.flags.writeable = False
            results = pose.process(image_rgb)
            image_rgb.flags.writeable = True
            image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
            the_color = calculate_color(similarity)
            if results.pose_landmarks:
                mp_drawing.draw_landmarks(
                    image_bgr,
                    results.pose_landmarks,
                    mp_pose.POSE_CONNECTIONS,
                    mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=4, circle_radius=2),
                    mp_drawing.DrawingSpec(color=(the_color), thickness=4, circle_radius=2),
                )

                landmarks_from_webcam = []
                for lm in results.pose_landmarks.landmark:
                    landmarks_from_webcam.append([lm.x, lm.y])

                similarity = calculate_similarity(landmarks_from_json, landmarks_from_webcam)
                cv2.putText(image_bgr, f'Similarity: {similarity:.2f}%', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

            ret, buffer = cv2.imencode('.jpg', image_bgr)
            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():
    image_name = request.args.get('image_name', 'Camel')
    previous = None
    next = None
    load_image_and_landmarks(image_name)

    current_index = 0
    for index, data in enumerate(all_data):
        if data['name'] == image_name:
            current_index = index
            break

    if current_index == 0:
        previous = all_data[-1]['name']
    else:
        previous = all_data[current_index - 1]['name']

    if current_index == len(all_data) - 1:
        next = all_data[0]['name']
    else:
        next = all_data[current_index + 1]['name']
    
    annotated_image = draw_landmarks(image, the_landmarks)
    _, buffer = cv2.imencode('.jpg', annotated_image)
    img_str = base64.b64encode(buffer).decode('utf-8')
    return render_template('index2.html', img_str=img_str, previous=previous, next=next)

@app.route('/video_feed')
def video_feed():
    return Response(generate_frames(), mimetype='multipart/x-mixed-replace; boundary=frame')

@app.route('/getdata', methods=['GET'])
def getdata():
    return jsonify(all_data)

@app.route('/similarity', methods=['GET'])
def get_similarity():
    global similarity
    return {'similarity': similarity, 'data': dataset}

@app.route('/pose_dataset')
def pose_dataset():
    load_image_and_landmarks("Camel")
    return render_template('pose_dataset.html', data=all_data)

@app.route('/show_image')
def show_image():
    image_path = request.args.get('image_path')
    if image_path and os.path.exists(image_path):
        return send_file(image_path, mimetype='image/jpeg')
    else:
        return "Image not found", 404

if __name__ == '__main__':
    socketio.run(app, debug=True)