import cv2
import mediapipe as mp
import numpy as np
from datetime import datetime
import os
import gradio as gr

# Create folder for dangerous frames
os.makedirs("dangerous_frames", exist_ok=True)

dangerous_move_count = 0

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

HOLD_NECK_THRESHOLD = 0.1  # Wrist-to-neck distance threshold for hold neck


# ---------------- Detection Logic ----------------
def detect_dangerous_move(landmarks):
    left_wrist = landmarks[15]
    right_wrist = landmarks[16]
    left_elbow = landmarks[13]
    right_elbow = landmarks[14]
    left_shoulder = landmarks[11]
    right_shoulder = landmarks[12]
    left_hip = landmarks[23]
    right_hip = landmarks[24]
    nose = landmarks[0]
    left_ankle = landmarks[27]
    right_ankle = landmarks[28]

    # Neck midpoint
    neck_x = (left_shoulder.x + right_shoulder.x) / 2
    neck_y = (left_shoulder.y + right_shoulder.y) / 2

    # Distances from wrists to neck
    left_wrist_to_neck = np.sqrt((left_wrist.x - neck_x) ** 2 + (left_wrist.y - neck_y) ** 2)
    right_wrist_to_neck = np.sqrt((right_wrist.x - neck_x) ** 2 + (right_wrist.y - neck_y) ** 2)

    # Punch
    punch_left = left_wrist.y < left_shoulder.y and left_elbow.y > left_wrist.y
    punch_right = right_wrist.y < right_shoulder.y and right_elbow.y > right_wrist.y

    # Kick
    kick_left = left_ankle.y < left_hip.y
    kick_right = right_ankle.y < right_hip.y

    # Block
    block_left = left_elbow.y < left_shoulder.y and abs(left_wrist.x - nose.x) < 0.15
    block_right = right_elbow.y < right_shoulder.y and abs(right_wrist.x - nose.x) < 0.15
    blocking = block_left and block_right

    # Hold neck
    hold_neck = (left_wrist_to_neck < HOLD_NECK_THRESHOLD) or (right_wrist_to_neck < HOLD_NECK_THRESHOLD)

    # Weapon raise
    weapon_left = (left_wrist.y < left_shoulder.y - 0.1) and (left_wrist.z < left_shoulder.z)
    weapon_right = (right_wrist.y < right_shoulder.y - 0.1) and (right_wrist.z < right_shoulder.z)

    if hold_neck:
        return "Hold Neck (Dangerous)", True
    elif punch_left or punch_right:
        return "Punching (Dangerous)", True
    elif kick_left or kick_right:
        return "Kicking (Dangerous)", True
    elif blocking:
        return "Blocking (Defensive)", False
    elif weapon_left or weapon_right:
        return "Weapon Raised (Dangerous)", True
    else:
        return "Normal", False


# ---------------- Main Gradio Function ----------------
def process_video(video_path):
    global dangerous_move_count
    dangerous_move_count = 0

    cap = cv2.VideoCapture(video_path)
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out = cv2.VideoWriter("output.mp4", fourcc, 20.0, (640, 480))

    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

            img_h, img_w, _ = frame.shape
            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            pose_results = pose.process(rgb)

            gesture = "No Person"
            danger_move = False
            landmarks = None

            if pose_results.pose_landmarks:
                landmarks = pose_results.pose_landmarks.landmark
                gesture, danger_move = detect_dangerous_move(landmarks)
                mp_drawing.draw_landmarks(frame, pose_results.pose_landmarks, mp_pose.POSE_CONNECTIONS)

            if landmarks:
                xs = [int(lm.x * img_w) for lm in landmarks]
                ys = [int(lm.y * img_h) for lm in landmarks]
                x_min, x_max = min(xs), max(xs)
                y_min, y_max = min(ys), max(ys)

                box_color = (0, 0, 255) if danger_move else (0, 255, 0)
                cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), box_color, 3)
                cv2.putText(frame, gesture, (x_min, y_min - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.9, box_color, 2)

                if danger_move:
                    dangerous_move_count += 1
                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
                    cv2.imwrite(f"dangerous_frames/danger_{timestamp}.jpg", frame)

            cv2.putText(frame, f"Dangerous Moves: {dangerous_move_count}", (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

            out.write(cv2.resize(frame, (640, 480)))

        cap.release()
        out.release()

    return "output.mp4"


# ---------------- Gradio Interface ----------------
demo = gr.Interface(
    fn=process_video,
    inputs=gr.Video(),       # <-- FIXED (removed type="filepath")
    outputs=gr.Video(),
    title="⚠️ Action Recognition System (Danger Detection)",
    description="Upload a video and the system will detect dangerous actions like punching, kicking, hold-neck, blocking, and weapon raise."
)

if __name__ == "__main__":
    demo.launch()