import os
import tempfile
import cv2
import numpy as np
import gradio as gr
import supervision as sv
from ultralytics import YOLO

model = YOLO("yolo11n.pt")
CLASS_NAMES_DICT = model.model.names

SELECTED_CLASS_NAMES = ['car', 'bus', 'truck', 'motorcycle']
SELECTED_CLASS_IDS = [
    {value: key for key, value in CLASS_NAMES_DICT.items()}[name]
    for name in SELECTED_CLASS_NAMES
]

box_annotator = sv.BoxAnnotator(thickness=4)
label_annotator = sv.LabelAnnotator(text_thickness=2, text_scale=1.5, text_color=sv.Color.BLACK)
trace_annotator = sv.TraceAnnotator(thickness=4, trace_length=50)
line_zone_annotator = sv.LineZoneAnnotator(
    thickness=4, text_thickness=2, text_scale=2,
    display_in_count=False, display_out_count=False
)


def process_video(video_path, orientation="Ngang"):
    """Xử lý video: nhận dạng phương tiện và trả về video đã annotate."""
    if video_path is None:
        return None

    if isinstance(video_path, dict):
        video_path = video_path.get("path", video_path)

    video_info = sv.VideoInfo.from_video_path(video_path)
    w, h = video_info.width, video_info.height

    ori = str(orientation).strip().lower()
    is_horizontal = ori.startswith("n") or ori.startswith("h")
    if is_horizontal:
        line_pos = int(h * 0.5)
        line_start = sv.Point(int(w * 0.005), line_pos)
        line_end = sv.Point(int(w * 0.995), line_pos)
    else:
        line_pos = int(w * 0.5)
        line_start = sv.Point(line_pos, int(h * 0.005))
        line_end = sv.Point(line_pos, int(h * 0.995))

    line_zone = sv.LineZone(start=line_start, end=line_end)
    byte_tracker = sv.ByteTrack(
        track_activation_threshold=0.25,
        lost_track_buffer=30,
        minimum_matching_threshold=0.8,
        frame_rate=video_info.fps or 30,
        minimum_consecutive_frames=3
    )
    byte_tracker.reset()

    previous_positions = {}
    class_counts = {name: 0 for name in SELECTED_CLASS_NAMES}
    crossed_ids = set()

    def callback(frame: np.ndarray, index: int) -> np.ndarray:
        nonlocal previous_positions, class_counts, crossed_ids

        results = model(frame, verbose=False)[0]
        detections = sv.Detections.from_ultralytics(results)
        detections = detections[np.isin(detections.class_id, SELECTED_CLASS_IDS)]
        detections = byte_tracker.update_with_detections(detections)

        if detections.tracker_id is not None:
            xyxy = detections.xyxy
            for i in range(len(detections)):
                tid = int(detections.tracker_id[i])
                cls_id = int(detections.class_id[i])
                cls_name = CLASS_NAMES_DICT[cls_id]
                cx = (xyxy[i, 0] + xyxy[i, 2]) / 2
                cy = (xyxy[i, 1] + xyxy[i, 3]) / 2

                curr_coord = cy if is_horizontal else cx
                if tid in previous_positions:
                    prev_coord = previous_positions[tid]
                    if prev_coord < line_pos and curr_coord > line_pos and (tid, 'pos') not in crossed_ids:
                        crossed_ids.add((tid, 'pos'))
                        class_counts[cls_name] = class_counts.get(cls_name, 0) + 1
                    elif prev_coord > line_pos and curr_coord < line_pos and (tid, 'neg') not in crossed_ids:
                        crossed_ids.add((tid, 'neg'))
                        class_counts[cls_name] = class_counts.get(cls_name, 0) + 1
                previous_positions[tid] = curr_coord

        labels = [
            f"#{tid} {CLASS_NAMES_DICT[cid]} {conf:0.2f}"
            for conf, cid, tid in zip(
                detections.confidence, detections.class_id, detections.tracker_id
            )
        ]

        annotator_frame = frame.copy()
        annotator_frame = trace_annotator.annotate(scene=annotator_frame, detections=detections)
        annotator_frame = box_annotator.annotate(scene=annotator_frame, detections=detections)
        annotator_frame = label_annotator.annotate(scene=annotator_frame, detections=detections, labels=labels)

        line_zone.trigger(detections)
        annotator_frame = line_zone_annotator.annotate(annotator_frame, line_counter=line_zone)

        fh, fw, _ = annotator_frame.shape
        box_w, box_h = 280, 50 + len(SELECTED_CLASS_NAMES) * 28
        x0, y0 = fw - box_w - 20, 20

        overlay = annotator_frame.copy()
        cv2.rectangle(overlay, (x0, y0), (x0 + box_w, y0 + box_h), (0, 0, 0), -1)
        annotator_frame = cv2.addWeighted(overlay, 0.6, annotator_frame, 0.4, 0)

        total = sum(class_counts.values())
        cv2.putText(annotator_frame, f'Total: {total}', (x0 + 10, y0 + 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
        for i, cls_name in enumerate(SELECTED_CLASS_NAMES):
            cnt = class_counts.get(cls_name, 0)
            cv2.putText(annotator_frame, f'{cls_name.capitalize()}: {cnt}',
                        (x0 + 10, y0 + 60 + i * 28),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)

        return annotator_frame

    output_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    sv.process_video(
        source_path=video_path,
        target_path=output_path,
        callback=callback
    )
    return output_path


with gr.Blocks(title="Object Detection", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🚗 Nhận dạng phương tiện (YOLOv8 + ByteTrack)")
    gr.Markdown("Tải lên video, hệ thống sẽ nhận dạng phương tiện trong video.")

    with gr.Row():
        video_input = gr.Video(label="Video đầu vào")
        video_output = gr.Video(label="Video đã xử lý")

    orientation_input = gr.Radio(
        choices=["Ngang", "Dọc"],
        value="Ngang",
        label="Vị trí line"
    )

    btn = gr.Button("▶️ Xử lý video")
    btn.click(fn=process_video, inputs=[video_input, orientation_input], outputs=video_output)

    gr.Markdown("""
    ### Author: Trần Hải Nam - 223332840
    """)

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