import gradio as gr
import cv2
from ultralytics import YOLO

# Initialize the YOLO model
model = YOLO("yolov8s.pt")

def process_video(video_path, analytics_type):
    cap = cv2.VideoCapture(video_path)
    assert cap.isOpened(), "Error reading video file"
    
    # Get video properties
    w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
    output_filename = f"{analytics_type}_output.avi"
    out = cv2.VideoWriter(output_filename, cv2.VideoWriter_fourcc(*"MJPG"), fps, (w, h))
    
    clswise_count = {}
    frame_count = 0

    while cap.isOpened():
        success, frame = cap.read()
        if success:
            frame_count += 1
            results = model.track(frame, persist=True, verbose=False)

            if results[0].boxes.id is not None:
                boxes = results[0].boxes.xyxy.cpu()
                clss = results[0].boxes.cls.cpu().tolist()

                for box, cls in zip(boxes, clss):
                    if model.names[int(cls)] in clswise_count:
                        clswise_count[model.names[int(cls)]] += 1
                    else:
                        clswise_count[model.names[int(cls)]] = 1

            # Perform simple analytics based on type
            if analytics_type == "line":
                # Display the number of detections on each frame (for example)
                cv2.putText(frame, f"Frame {frame_count}: Detections - {sum(clswise_count.values())}", 
                            (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, cv2.LINE_AA)
            elif analytics_type == "multiple_line":
                # Display classwise counts
                for i, (cls_name, count) in enumerate(clswise_count.items()):
                    cv2.putText(frame, f"{cls_name}: {count}", 
                                (10, 30 + (i + 1) * 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, cv2.LINE_AA)
            elif analytics_type == "pie":
                # Placeholder for pie chart (can implement later with matplotlib and overlay on frame)
                pass
            elif analytics_type == "area":
                # Placeholder for area graph (implement as needed)
                pass

            out.write(frame)

            clswise_count = {}  # Reset for next frame

        else:
            break

    cap.release()
    out.release()

    return output_filename  # Return the output video file


def gradio_app(video, analytics_type):
    # Save uploaded video locally
    video_path = video  # Gradio automatically returns the file path
    output_video = process_video(video_path, analytics_type)
    
    # Return processed video for display
    return output_video


# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# YOLO Video Processing App")
    
    with gr.Row():
        video_input = gr.Video(label="Upload Video")
        analytics_dropdown = gr.Dropdown(
            ["line", "multiple_line", "pie", "area"], 
            label="Select Analytics Type", 
            value="line"
        )
    
    output_video = gr.Video(label="Processed Output")
    
    # Button to start processing
    submit_btn = gr.Button("Process Video")

    # Define the output when the button is clicked
    submit_btn.click(gradio_app, inputs=[video_input, analytics_dropdown], outputs=output_video)

# Launch the Gradio app with a public link
demo.launch(share=True)