import gradio as gr
import cv2
from ultralytics import YOLO
import os
from moviepy.editor import VideoFileClip

# 1. Khởi tạo mô hình
model = YOLO('yolov8n.pt')

def process_video(video_path):
    if not video_path:
        return None

    # --- CẤU HÌNH BIẾN ĐẾM (Dựa trên code gốc của bạn) ---
    vehicle_classes = ['car', 'motorcycle', 'bus', 'truck']
    vehicle_class_ids = [2, 3, 5, 7] 
    vehicle_count = {cls: 0 for cls in vehicle_classes}
    
    track_history = {}  
    counted_ids = set() 

    # Đọc video đầu vào
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    # Vạch đếm ở 60% chiều cao màn hình
    line_y = int(height * 0.6) 

    # File tạm trung gian (định dạng mp4v)
    temp_output = 'temp_raw.mp4'
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(temp_output, fourcc, fps, (width, height))

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

        # Thực hiện Tracking
        results = model.track(frame, classes=vehicle_class_ids, persist=True, verbose=False)

        # Vẽ vạch đếm màu đỏ
        cv2.line(frame, (0, line_y), (width, line_y), (0, 0, 255), 3)
        cv2.putText(frame, "Counting Line", (10, line_y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)

        if results[0].boxes.id is not None:
            boxes = results[0].boxes
            ids = boxes.id.cpu().numpy().astype(int)
            classes = boxes.cls.cpu().numpy().astype(int)

            for (x1, y1, x2, y2, obj_id, cls_id) in zip(
                boxes.xyxy[:,0], boxes.xyxy[:,1], boxes.xyxy[:,2], boxes.xyxy[:,3], ids, classes
            ):
                label = model.names[cls_id]
                cx = int((x1 + x2) / 2)
                cy = int((y1 + y2) / 2)

                # Vẽ khung và tâm xe
                cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (255, 255, 0), 2)
                cv2.circle(frame, (cx, cy), 5, (0, 255, 0), -1) 

                # Logic đếm qua vạch
                if obj_id in track_history and obj_id not in counted_ids:
                    prev_cy = track_history[obj_id]
                    if (prev_cy < line_y and cy >= line_y) or (prev_cy > line_y and cy <= line_y):
                        vehicle_count[label] += 1
                        counted_ids.add(obj_id) 
                        cv2.line(frame, (0, line_y), (width, line_y), (0, 255, 0), 5) 
                
                track_history[obj_id] = cy

        # --- BẢNG THỐNG KÊ (Thiết kế gốc của bạn) ---
        total = sum(vehicle_count.values())
        overlay = frame.copy()
        box_w, box_h = 250, 160
        x0, y0 = width - box_w - 20, 20
        cv2.rectangle(overlay, (x0, y0), (x0 + box_w, y0 + box_h), (0, 0, 0), -1)
        frame = cv2.addWeighted(overlay, 0.6, frame, 0.4, 0)

        cv2.putText(frame, f'Total: {total}', (x0 + 10, y0 + 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)

        for i, (cls, cnt) in enumerate(vehicle_count.items()):
            cv2.putText(frame, f'{cls.capitalize()}: {cnt}',
                        (x0 + 10, y0 + 60 + i * 25),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)

        out.write(frame)

    cap.release()
    out.release()

    # --- CHUYỂN ĐỔI SANG CHUẨN H.264 ĐỂ XEM TRỰC TIẾP ---
    final_output = 'result_display.mp4'
    try:
        # Chuyển đổi định dạng để trình duyệt web có thể Play trực tiếp
        clip = VideoFileClip(temp_output)
        clip.write_videofile(final_output, codec="libx264", audio=False)
        clip.close()
        if os.path.exists(temp_output):
            os.remove(temp_output)
    except Exception as e:
        print(f"Lỗi encode video: {e}")
        return temp_output
    
    return final_output

# --- GIAO DIỆN GRADIO ---
interface = gr.Interface(
    fn=process_video,
    inputs=gr.Video(label="Tải video của bạn lên"),
    outputs=gr.Video(label="Kết quả đếm xe (Xem trực tiếp)"),
    title="🚗 Hệ thống Đếm Xe Thông Minh (YOLOv8)",
    description="Nhận diện và đếm số lượng xe qua vạch. Kết quả xem được trực tiếp trên web.",
)

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