import cv2
import numpy as np
import pandas as pd
import math
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
import os

# ============================================
# USER SETTINGS
# ============================================

VIDEO_PATH = r"C:\Users\DURGA\Downloads\ad_28.mp4"
OUTPUT_FOLDER = r"C:\Users\DURGA\Downloads"

FPS = 30
MIN_CONTOUR_AREA = 60
STOP_DISTANCE = 3

# ============================================
# TRACKING FUNCTION
# ============================================

def analyze_video(video_path):

    cap = cv2.VideoCapture(video_path)

    positions = []
    distances = []
    stops = 0
    turns = 0

    prev = None
    prev_dx = None

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        blur = cv2.GaussianBlur(gray, (5, 5), 0)

        _, thresh = cv2.threshold(
            blur, 0, 255,
            cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU
        )

        kernel = np.ones((3, 3), np.uint8)
        thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
        thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)

        contours, _ = cv2.findContours(
            thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )

        centroids = []
        for c in contours:
            if cv2.contourArea(c) > MIN_CONTOUR_AREA:
                M = cv2.moments(c)
                if M["m00"] != 0:
                    cx = int(M["m10"] / M["m00"])
                    cy = int(M["m01"] / M["m00"])
                    centroids.append((cx, cy))

        if len(centroids) > 0:
            avg_x = int(np.mean([c[0] for c in centroids]))
            avg_y = int(np.mean([c[1] for c in centroids]))
            positions.append((avg_x, avg_y))

            if prev is not None:
                dx = avg_x - prev[0]
                dy = prev[1] - avg_y
                dist = math.sqrt(dx**2 + dy**2)
                distances.append(dist)

                if dist < STOP_DISTANCE:
                    stops += 1

                if prev_dx is not None and dx * prev_dx < 0:
                    turns += 1

                prev_dx = dx

            prev = (avg_x, avg_y)

    cap.release()

    # ============================================
    # BASIC METRICS
    # ============================================

    total_distance = sum(distances)
    mean_velocity = total_distance / (len(distances) / FPS + 1e-6)
    peak_velocity = max(distances) * FPS if distances else 0

    climbing_height = (
        positions[0][1] - positions[-1][1]
        if len(positions) > 1 else 0
    )

    traj_area = 0
    if len(positions) > 5:
        hull = cv2.convexHull(np.array(positions, dtype=np.int32))
        traj_area = cv2.contourArea(hull)

    # ============================================
    # FREECLIMBER-STYLE SLOPE CALCULATION
    # ============================================

    slope = 0

    if len(positions) > FPS:

        ys = np.array([p[1] for p in positions])
        ys = max(ys) - ys  # invert so upward positive

        frames = np.arange(len(ys))
        window_size = FPS  # 1-second window

        best_r2 = -1
        best_slope = 0
        best_start = 0

        for i in range(len(frames) - window_size):

            X = frames[i:i+window_size].reshape(-1, 1)
            y = ys[i:i+window_size]

            model = LinearRegression().fit(X, y)
            r2 = model.score(X, y)

            if r2 > best_r2:
                best_r2 = r2
                best_slope = model.coef_[0] * FPS
                best_start = i

        slope = best_slope

        # ============================================
        # PLOT SLOPE SEGMENT
        # ============================================

        plt.figure(figsize=(6, 4))
        plt.plot(frames / FPS, ys, label="Mean vertical position")

        best_end = best_start + window_size

        plt.plot(
            frames[best_start:best_end] / FPS,
            ys[best_start:best_end],
            color='red',
            linewidth=3,
            label="Most linear segment"
        )

        plt.xlabel("Time (seconds)")
        plt.ylabel("Climbing height (pixels)")
        plt.title(f"Slope = {slope:.2f} px/sec")
        plt.legend()
        plt.tight_layout()

        plot_path = os.path.join(OUTPUT_FOLDER, "slope_plot.png")
        plt.savefig(plot_path)
        plt.close()

    return {
        "climbing_height": climbing_height,
        "mean_velocity": mean_velocity,
        "peak_velocity": peak_velocity,
        "trajectory_length": total_distance,
        "trajectory_area": traj_area,
        "turn_count": turns,
        "stop_count": stops,
        "climbing_slope": slope
    }


# ============================================
# RUN ANALYSIS
# ============================================

results = analyze_video(VIDEO_PATH)

df = pd.DataFrame([results])

csv_path = os.path.join(OUTPUT_FOLDER, "analysis_output.csv")
df.to_csv(csv_path, index=False)

print("\nAnalysis complete.")
print(df)