import gradio as gr
import cv2
import numpy as np
import tensorflow as tf
from ultralytics import YOLO
import tempfile
import os
import time

# ──────────────────────────────────────────
# LOAD MODELS
# ──────────────────────────────────────────
print("Loading YOLO...")
yolo_model = YOLO("models/yolo.pt")

print("Loading EfficientNet-B0...")
# Loads your newly upgraded model architecture
resnet_model = tf.keras.models.load_model("models/efficientnetb0.keras", compile=False)

# Explicitly hardcoding your optimal benchmark evaluation threshold for EfficientNet-B0
RESNET_THRESHOLD = 0.35

CONF_THRESHOLD  = 0.20
HEAD_CROP_RATIO = 0.80   

# Warm up the model
dummy = np.zeros((1, 224, 224, 3), dtype="float32")
resnet_model.predict(dummy, verbose=0)

print(f"[✓] YOLO classes : {yolo_model.names}")
print(f"[✓] Optimal EffNet Threshold : {RESNET_THRESHOLD:.2f}")

VEHICLE_CLASSES = {"motorcycle", "bike", "car", "truck", "bus", "vehicle", "auto", "scooter"}

# ──────────────────────────────────────────
# HELPERS
# ──────────────────────────────────────────
def preprocess_crop(crop_bgr):
    # Matches the exact preprocessing dimensions 
    crop = cv2.resize(crop_bgr, (224, 224))
    crop = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB)
    crop = crop.astype("float32")
    return np.expand_dims(crop, axis=0)

def batch_predict(crops):
    if not crops:
        return []
    batch = np.concatenate(crops, axis=0)
    preds = resnet_model.predict(batch, verbose=0)
    # Check if binary (sigmoid) vs categorical (softmax)
    if preds.shape[-1] == 1:
        return preds.flatten().tolist()
    else:
        # If model outputs 2 classes, extract probability belonging to 'no_helmet' (class 1)
        return preds[:, 1].tolist()

def draw_box(frame, x1, y1, x2, y2, label, color):
    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
    (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
    cv2.rectangle(frame, (x1, y1 - th - 10), (x1 + tw + 6, y1), color, -1)
    cv2.putText(frame, label, (x1 + 3, y1 - 4),
                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)

def draw_summary(frame, riders, helmets, no_helmets):
    pct     = int(helmets / max(riders, 1) * 100)
    summary = f"Riders:{riders}  Helmet:{helmets}  NoHelmet:{no_helmets}  {pct}%"
    bar_w   = min(len(summary) * 13, frame.shape[1])
    cv2.rectangle(frame, (0, 0), (bar_w, 36), (20, 20, 20), -1)
    cv2.putText(frame, summary, (8, 26),
                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)

def classify_by_name(cls_name):
    n = cls_name.lower().strip()
    if any(x in n for x in ["without", "no helmet", "nohelmet", "no-helmet"]):
        return "no_helmet"
    if any(x in n for x in ["with helmet", "helmet"]):
        return "helmet"
    return "person"

def process_frame(frame):
    results   = yolo_model(frame, conf=CONF_THRESHOLD, verbose=False)
    annotated = frame.copy()

    boxes_info = []   
    crops_list = []   

    for r in results:
        for box in r.boxes:
            cls_id   = int(box.cls[0])
            cls_name = yolo_model.names[cls_id].lower().strip()

            if cls_name in VEHICLE_CLASSES:
                continue

            x1, y1, x2, y2 = map(int, box.xyxy[0])
            x1 = max(0, x1); y1 = max(0, y1)
            x2 = min(frame.shape[1], x2)
            y2 = min(frame.shape[0], y2)

            if (x2 - x1) < 15 or (y2 - y1) < 15:
                continue

            mode = classify_by_name(cls_name)

            if mode in ("helmet", "no_helmet"):
                boxes_info.append((x1, y1, x2, y2, mode))
                crops_list.append(None)
            else:
                head_y2 = min(y1 + int((y2 - y1) * HEAD_CROP_RATIO), frame.shape[0])
                crop = frame[y1:head_y2, x1:x2]
                if crop.size == 0:
                    continue
                boxes_info.append((x1, y1, x2, head_y2, "efficientnet"))
                crops_list.append(preprocess_crop(crop))

    resnet_crops = [c for c in crops_list if c is not None]
    resnet_preds = batch_predict(resnet_crops)
    pred_iter    = iter(resnet_preds)

    riders = helmets = no_helmets = 0

    for (x1, y1, x2, y2, mode) in boxes_info:
        riders += 1

        if mode == "helmet":
            label = "Helmet"
            conf  = 0.94
        elif mode == "no_helmet":
            label = "No Helmet"
            conf  = 0.94
        else:
            prob  = next(pred_iter)
            label = "No Helmet" if prob >= RESNET_THRESHOLD else "Helmet"
            conf  = prob if label == "No Helmet" else (1 - prob)

        color = (0, 200, 0) if label == "Helmet" else (0, 0, 220)

        if label == "Helmet":
            helmets += 1
        else:
            no_helmets += 1

        draw_box(annotated, x1, y1, x2, y2, f"{label} {conf:.2f}", color)

    draw_summary(annotated, riders, helmets, no_helmets)
    return annotated, riders, helmets, no_helmets

# ──────────────────────────────────────────
# INTERFACE IMPLEMENTATIONS
# ──────────────────────────────────────────
def predict_image(image_rgb):
    if image_rgb is None:
        return None, "No image uploaded."
    frame      = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
    result, riders, helmets, no_helmets = process_frame(frame)
    result_rgb = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
    stats = (
        f"Total Riders : {riders}\n"
        f"With Helmet  : {helmets}\n"
        f"No Helmet    : {no_helmets}\n"
        f"Compliance   : {int(helmets / max(riders,1) * 100)}%"
    )
    return result_rgb, stats

def predict_video(video_path, skip_frames, max_minutes, progress=gr.Progress()):
    if video_path is None:
        return None, "No video uploaded."

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return None, "Cannot open video."

    width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps    = cap.get(cv2.CAP_PROP_FPS) or 25.0
    total  = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    scale = 1.0
    if width > 640:
        scale  = 640.0 / width
        width  = 640
        height = int(height * scale)

    max_frames = int(max_minutes * 60 * fps)
    total      = min(total, max_frames)

    out_path = tempfile.mktemp(suffix=".mp4")
    writer   = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height))

    skip           = max(1, int(skip_frames))
    frame_idx      = 0
    proc_count     = 0
    last_annotated = None

    total_helmets  = 0
    total_no_helmets = 0
    frames_with_riders = 0
    start_time     = time.time()

    progress(0, desc="Starting...")

    while frame_idx < total:
        ret, frame = cap.read()
        if not ret:
            break

        frame_idx += 1
        if scale < 1.0:
            frame = cv2.resize(frame, (width, height))

        if frame_idx % skip == 0:
            annotated, r, h, nh = process_frame(frame)
            last_annotated       = annotated
            proc_count          += 1

            if r > 0:
                frames_with_riders += 1
                total_helmets      += h
                total_no_helmets   += nh
        else:
            annotated = last_annotated if last_annotated is not None else frame

        writer.write(annotated)

        if frame_idx % 10 == 0:
            elapsed = time.time() - start_time
            spd     = proc_count / max(elapsed, 1)
            progress(frame_idx / total, desc=f"Frame {frame_idx}/{total} | {spd:.1f} fps")

    cap.release()
    writer.release()

    elapsed  = time.time() - start_time
    avg_h    = total_helmets    // max(frames_with_riders, 1)
    avg_nh   = total_no_helmets // max(frames_with_riders, 1)
    avg_r    = (avg_h + avg_nh)

    stats = (
        f"Total frames       : {frame_idx}\n"
        f"Frames processed   : {proc_count}\n"
        f"Processing time    : {elapsed:.1f}s\n"
        f"Speed              : {proc_count/max(elapsed,1):.1f} fps\n"
        f"Avg Helmet/frame   : {avg_h}\n"
        f"Avg No Helmet      : {avg_nh}\n"
        f"Compliance rate    : {int(avg_h/max(avg_r,1)*100)}%"
    )
    return out_path, stats

# ──────────────────────────────────────────
# GRADIO BLOCK UI BUILD
# ──────────────────────────────────────────
with gr.Blocks(title="Smart Helmet Detection", theme=gr.themes.Soft(), css="footer{display:none!important} .mono{font-family:monospace}") as demo:
    gr.Markdown("# 🏍️ Smart Helmet Detection System")
    gr.Markdown("**YOLOv8 + EfficientNet-B0** — Upload an image or video")

    with gr.Tabs():
        with gr.Tab("📷 Image Detection"):
            with gr.Row():
                with gr.Column():
                    img_input  = gr.Image(label="Upload Image", type="numpy", sources=["upload","clipboard"])
                    img_button = gr.Button("🔍 Detect Helmets", variant="primary", size="lg")
                with gr.Column():
                    img_output = gr.Image(label="Detection Result")
                    img_stats  = gr.Textbox(label="Statistics", lines=5, elem_classes="mono")
            img_button.click(predict_image, inputs=img_input, outputs=[img_output, img_stats])

        with gr.Tab("🎥 Video Detection"):
            with gr.Row():
                with gr.Column():
                    vid_input = gr.Video(label="Upload Video", sources=["upload"])
                    with gr.Row():
                        skip_slider = gr.Slider(minimum=1, maximum=15, value=5, step=1, label="Skip frames")
                        max_min_slider = gr.Slider(minimum=0.5, maximum=3.0, value=0.5, step=0.5, label="Max minutes")
                    vid_button = gr.Button("🎬 Process Video", variant="primary", size="lg")
                with gr.Column():
                    vid_output = gr.Video(label="Result Video")
                    vid_stats  = gr.Textbox(label="Statistics", lines=10, elem_classes="mono")
            vid_button.click(predict_video, inputs=[vid_input, skip_slider, max_min_slider], outputs=[vid_output, vid_stats])

demo.queue(max_size=5).launch()