import cv2
import numpy as np
import torch
from ultralytics import YOLO
from sort import Sort
import gradio as gr

# Load YOLOv12x model
MODEL_PATH = "yolov12x.pt"
model = YOLO(MODEL_PATH)

# COCO dataset class ID for truck
TRUCK_CLASS_ID = 7  # "truck"

# Initialize SORT tracker
tracker = Sort()

def count_unique_trucks(video_path, video_type, confidence_threshold, distance_threshold, frame_skip_seconds):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return "Error: Unable to open video file."

    unique_truck_ids = set()
    truck_history = {}

    # Get FPS of the video
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    frame_skip = max(1, int(fps * frame_skip_seconds))  # Dynamic frame skipping

    frame_count = 0

    while True:
        ret, frame = cap.read()
        if not ret:
            break  # End of video

        frame_count += 1
        if frame_count % frame_skip != 0:
            continue  # Skip frames dynamically

        # Run YOLOv12x inference
        results = model(frame, verbose=False)

        detections = []
        for result in results:
            for box in result.boxes:
                class_id = int(box.cls.item())  # Get class ID
                confidence = float(box.conf.item())  # Get confidence score

                # Track only trucks
                if class_id == TRUCK_CLASS_ID and confidence > confidence_threshold:
                    x1, y1, x2, y2 = map(int, box.xyxy[0])  # Get bounding box
                    detections.append([x1, y1, x2, y2, confidence])

        if len(detections) > 0:
            detections = np.array(detections)
            tracked_objects = tracker.update(detections)

            for obj in tracked_objects:
                truck_id = int(obj[4])  # Unique ID assigned by SORT
                x1, y1, x2, y2 = obj[:4]  # Get the bounding box coordinates

                truck_center = (x1 + x2) / 2, (y1 + y2) / 2  # Calculate the center of the truck

                # If truck is already in history, check the movement distance
                if truck_id in truck_history:
                    last_position = truck_history[truck_id]["position"]
                    distance = np.linalg.norm(np.array(truck_center) - np.array(last_position))

                    if distance > distance_threshold:
                        # If the truck moved significantly, count as new
                        unique_truck_ids.add(truck_id)

                else:
                    # If truck is not in history, add it
                    truck_history[truck_id] = {
                        "frame_count": frame_count,
                        "position": truck_center
                    }
                    unique_truck_ids.add(truck_id)

    cap.release()
    return {"Total Unique Trucks": len(unique_truck_ids)}

# Gradio UI function
def analyze_video(video_file, video_type, confidence, distance, frame_skip):
    return count_unique_trucks(video_file, video_type, confidence, distance, frame_skip)

# Define Gradio interface
iface = gr.Interface(
    fn=analyze_video,
    inputs=[
        gr.Video(label="Upload Video"),
        gr.Radio(["drone", "fixed"], label="Video Type"),
        gr.Slider(0.3, 0.9, 0.5, label="Confidence Threshold"),
        gr.Slider(10, 100, 50, label="Distance Threshold"),
        gr.Slider(1, 10, 2, label="Frame Skip (Seconds)"),
    ],
    outputs=gr.JSON(label="Analysis Result"),
    title="YOLOv12x Dynamic Truck Counter",
    description="Upload a video, adjust parameters, and analyze unique trucks."
)

# Launch the Gradio app
if __name__ == "__main__":
    iface.launch()