import gradio as gr
import cv2
import numpy as np
import tempfile
import os

def detect_and_predict(video):
    # Save uploaded video
    temp_video_path = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4').name
    with open(temp_video_path, "wb") as f:
        f.write(video.read())

    cap = cv2.VideoCapture(temp_video_path)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out_path = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4').name
    out = cv2.VideoWriter(out_path, fourcc, 20.0, (width, height))

    ball_color_lower = np.array([5, 50, 50])     # Orange/red lower
    ball_color_upper = np.array([15, 255, 255])  # Orange/red upper
    trajectory_points = []

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

        blurred = cv2.GaussianBlur(frame, (11, 11), 0)
        hsv = cv2.cvtColor(blurred, cv2.COLOR_BGR2HSV)
        mask = cv2.inRange(hsv, ball_color_lower, ball_color_upper)
        mask = cv2.erode(mask, None, iterations=2)
        mask = cv2.dilate(mask, None, iterations=2)

        contours, _ = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        if contours:
            c = max(contours, key=cv2.contourArea)
            ((x, y), radius) = cv2.minEnclosingCircle(c)

            if radius > 3:
                trajectory_points.append((int(x), int(y)))
                cv2.circle(frame, (int(x), int(y)), int(radius), (0, 0, 255), 2)

        # Draw trajectory
        for i in range(1, len(trajectory_points)):
            cv2.line(frame, trajectory_points[i - 1], trajectory_points[i], (255, 0, 0), 2)

        # Draw stumps line (for simplicity, fixed zone)
        cv2.rectangle(frame, (width // 2 - 20, height - 200), (width // 2 + 20, height - 50), (0, 255, 255), 2)

        out.write(frame)

    cap.release()
    out.release()

    return out_path

iface = gr.Interface(fn=detect_and_predict,
                     inputs=gr.Video(label="Upload Bowling Video"),
                     outputs=gr.Video(label="Ball Tracking Result"),
                     title="DRS Ball Tracker",
                     description="Detect and visualize ball trajectory for LBW simulation.")

iface.launch()