import numpy as np
import supervision as sv
from ultralytics import YOLO
import gradio as gr
import torch
import spaces
import os

tracker = sv.ByteTrack()
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
trace_annotator = sv.TraceAnnotator()

device = "cuda" if torch.cuda.is_available() else "cpu"
model = YOLO("yolov8n.pt").to(device)

FRAME_INTERVAL = 35
LOW_RES_SIZE = (640,640)
MID_RES_SIZE = (1280,1280)
HIGH_RES_SIZE =(1088,1920)
MAX_PERSONS = 0

auth_users = [(os.getenv('USERNAME'), os.getenv('PASSWORD'))]

@spaces.GPU
def process_video(frame: np.ndarray, _: int) -> np.ndarray:
    global MAX_PERSONS
    
    results = model(frame, imgsz=HIGH_RES_SIZE)[0]
    detections = sv.Detections.from_ultralytics(results)
    person_detections = detections[detections.class_id == 0]
    
    current_person_count = len(person_detections)
    if current_person_count > MAX_PERSONS:
        MAX_PERSONS = current_person_count
    
    print(f'Personas detectadas: {current_person_count}, Máximo: {MAX_PERSONS}')
    
    tracked_persons = tracker.update_with_detections(person_detections)
    
    labels = [
        f"#{tracker_id} {results.names[class_id]}"
        for class_id, tracker_id
        in zip(tracked_persons.class_id, tracked_persons.tracker_id)
    ]
    
    annotated_frame = box_annotator.annotate(frame.copy(), detections=tracked_persons)
    annotated_frame = label_annotator.annotate(annotated_frame, detections=tracked_persons, labels=labels)
    
    return trace_annotator.annotate(annotated_frame, detections=tracked_persons)

def upload_video(video_path):
    global MAX_PERSONS
    MAX_PERSONS = 0  # Reiniciar el contador para cada video

    output_video_path = "output_video.mp4"

    if os.path.exists(output_video_path):
        os.remove(output_video_path)

    sv.process_video(source_path=video_path, target_path=output_video_path, callback=process_video)
    
    return output_video_path, MAX_PERSONS

with gr.Blocks() as demo:
    gr.Markdown("# Aegis Air Demo")
    with gr.Row():
        video_input = gr.Video(label="Sube tu video aquí")
        
            
        video_output = gr.Video(label="Video Anotado")
        
    persons_output = gr.Textbox(label="Total de Personas Detectadas", interactive=False)    
    def process_video_gradio(video):
        if video is None:
            return None, "No se ha subido ningún video"
            
        processed_video, total_persons = upload_video(video)
        return processed_video, f"Máximo número de personas detectadas: {total_persons}"
        
    submit_button = gr.Button("Procesar")
    submit_button.click(process_video_gradio, inputs=video_input, outputs=[video_output, persons_output])
    
demo.launch(auth=auth_users)