import cv2
import os
import numpy as np
import matplotlib
matplotlib.use('Agg')  # Use non-GUI backend for Flask
import matplotlib.pyplot as plt
from sklearn.metrics import mean_squared_error
from scipy.interpolate import interp1d
import mediapipe as mp

def analyze_video(video_path, video_id, user_id):
    mp_pose = mp.solutions.pose
    pose = mp_pose.Pose(static_image_mode=False, min_detection_confidence=0.5, min_tracking_confidence=0.5)

    cap = cv2.VideoCapture(video_path)
    frames = []
    sacrum_positions = []
    pitch_values = []

    frame_count = 0
    max_frames = int(cap.get(cv2.CAP_PROP_FPS)) * 10

    while cap.isOpened() and frame_count < max_frames:
        ret, frame = cap.read()
        if not ret:
            break
        frame_count += 1

        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        results = pose.process(frame_rgb)

        if results.pose_landmarks:
            lm = results.pose_landmarks.landmark
            left_hip = lm[mp_pose.PoseLandmark.LEFT_HIP.value]
            right_hip = lm[mp_pose.PoseLandmark.RIGHT_HIP.value]
            left_shoulder = lm[mp_pose.PoseLandmark.LEFT_SHOULDER.value]
            right_shoulder = lm[mp_pose.PoseLandmark.RIGHT_SHOULDER.value]

            visibility = (
                left_hip.visibility + right_hip.visibility +
                left_shoulder.visibility + right_shoulder.visibility
            ) / 4

            if visibility > 0.5:
                sacrum_z = (left_hip.z + right_hip.z) / 2
                pitch = (left_shoulder.y + right_shoulder.y)/2 - (left_hip.y + right_hip.y)/2

                sacrum_positions.append(sacrum_z)
                pitch_values.append(pitch)
                frames.append(frame)

    cap.release()

    if not sacrum_positions or len(sacrum_positions) < 5:
        raise Exception("Insufficient motion data")

    def resample_to_150(arr):
        x = np.linspace(0, 1, len(arr))
        f = interp1d(x, arr, kind='linear')
        return f(np.linspace(0, 1, 150))

    surf_z = resample_to_150(sacrum_positions)
    surf_pitch = resample_to_150(pitch_values)

    best_z = np.concatenate([
        np.linspace(0.1, 0.2, 30),
        np.linspace(0.2, 0.6, 30),
        np.linspace(0.6, 0.65, 20),
        np.linspace(0.65, 0.5, 20),
        np.linspace(0.5, 0.55, 20),
        np.linspace(0.55, 0.6, 30)
    ])
    best_pitch = np.concatenate([
        np.linspace(0.0, 0.3, 30),
        np.linspace(0.3, 0.8, 30),
        np.linspace(0.8, 0.6, 20),
        np.linspace(0.6, 0.5, 20),
        np.linspace(0.5, 0.4, 20),
        np.linspace(0.4, 0.4, 30)
    ])

    def normalize_range(arr):
        arr = np.array(arr)
        return (arr - arr.min()) / (arr.max() - arr.min() + 1e-6)

    surf_z_norm = normalize_range(surf_z)
    surf_pitch_norm = normalize_range(surf_pitch)
    best_z_norm = normalize_range(best_z)
    best_pitch_norm = normalize_range(best_pitch)

    stages = ['Start', 'Push-Up', 'Peak', 'Foot Contact', 'Dip', 'End']
    boundaries = [0, 30, 60, 80, 100, 120, 150]

    fig, axs = plt.subplots(2, 1, figsize=(10, 6), sharex=True)

    axs[0].plot(surf_z_norm, label="Surfer Sacrum Z")
    axs[0].plot(best_z_norm, label="Best Practice", linestyle="--")
    axs[0].set_title("Vertical Displacement")
    axs[0].legend()

    axs[1].plot(surf_pitch_norm, label="Surfer Pitch")
    axs[1].plot(best_pitch_norm, label="Best Practice", linestyle="--")
    axs[1].set_title("Torso Pitch")
    axs[1].legend()

    for i in range(6):
        axs[0].axvspan(boundaries[i], boundaries[i+1], alpha=0.1)
        axs[1].axvspan(boundaries[i], boundaries[i+1], alpha=0.1)

    plt.xlabel("Normalized Time (0–150 points)")
    plt.tight_layout()

    # Save figure to /tmp/results/{user_id}/
    output_folder = '/tmp/results'
    os.makedirs(output_folder, exist_ok=True)

    output_file = os.path.join(output_folder, f"{video_id}_popup_analysis.png")
    plot_filename = f"{video_id}_popup_analysis.png"
    
    plt.savefig(output_file)
    plt.close()

    print(f"✅ Saved result to: {output_file}")
    print(f"✅ Returning filename: {video_id}_popup_analysis.png")
    
    # Save 2–3 key frames as examples
    frame_filenames = []
    for i, frame in enumerate(frames[:3]):  # Adjust how many you want
        frame_filename = f"frame_{i}.jpg"
        frame_path = os.path.join(output_folder, frame_filename)
        cv2.imwrite(frame_path, frame)
        frame_filenames.append(frame_filename)

    # Log output
    print(f"✅ Saved plot to: {output_file}")
    print(f"📤 Returning frames: {frame_filenames}")

    # plot_filename = f"{video_id}_popup_analysis.png"
    # frame_filenames = [f"frame_0.jpg", "frame_1.jpg", ...]

    return {
        "plot": plot_filename,
        "frames": frame_filenames
    }