"""
HuggingFace Spaces — Gradio Demo
Mudflap / Black-sheet Detection  |  Models: YOLOv11x + YOLOv5x
"""

import os
import gradio as gr
import torch
import numpy as np
from pathlib import Path
from PIL import Image, ImageDraw, ImageFont
import time

os.environ.setdefault("YOLO_CONFIG_DIR", "/tmp/Ultralytics")

# ── Config ────────────────────────────────────────────────────
DEVICE      = "cuda" if torch.cuda.is_available() else "cpu"
MODEL_DIR   = Path("models")
CLASS_NAMES = ["mudflap"]
COLORS      = {"YOLOv11x": (0, 200, 100), "YOLOv5x": (255, 140, 0)}

# ── Model loading ─────────────────────────────────────────────
def load_v11():
    from ultralytics import YOLO
    m = YOLO(str(MODEL_DIR / "yolov11x_mudflap_best.pt"))
    m.to(DEVICE)
    print("[✓] YOLOv11x loaded.")
    return m

def load_v5():
    m = torch.hub.load(
        "ultralytics/yolov5", "custom",
        path=str(MODEL_DIR / "yolov5x_mudflap_best.pt"),
        force_reload=False, trust_repo=True, device=DEVICE,
    )
    m.conf = 0.60; m.iou = 0.45; m.max_det = 100
    print("[✓] YOLOv5x loaded.")
    return m

model_v11 = load_v11()
model_v5  = load_v5()

# ── Inference ─────────────────────────────────────────────────
def run_v11(img_np, conf_thr, iou_thr):
    results = model_v11.predict(source=img_np, conf=conf_thr, iou=iou_thr,
                                imgsz=640, device=DEVICE, verbose=False)
    boxes = []
    for r in results:
        for b in r.boxes:
            x1, y1, x2, y2 = b.xyxy[0].cpu().tolist()
            boxes.append((x1, y1, x2, y2, float(b.conf[0]), int(b.cls[0])))
    return boxes

def run_v5(img_np, conf_thr, iou_thr):
    model_v5.conf = conf_thr; model_v5.iou = iou_thr
    results = model_v5(img_np, size=640)
    return [(x[0], x[1], x[2], x[3], x[4], int(x[5]))
            for *x, _ in [r for r in results.xyxy[0].cpu().tolist()]]

def run_v5(img_np, conf_thr, iou_thr):
    model_v5.conf = conf_thr
    model_v5.iou  = iou_thr
    results = model_v5(img_np, size=640)
    boxes = []
    for *xyxy, conf, cls in results.xyxy[0].cpu().tolist():
        boxes.append((xyxy[0], xyxy[1], xyxy[2], xyxy[3], conf, int(cls)))
    return boxes

# ── Drawing ───────────────────────────────────────────────────
def draw_boxes(img_pil, detections, color, prefix):
    draw = ImageDraw.Draw(img_pil)
    try:
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 18)
    except Exception:
        font = ImageFont.load_default()
    for (x1, y1, x2, y2, conf, cls) in detections:
        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
        lbl = f"{prefix} {CLASS_NAMES[cls]} {conf:.2f}"
        r, g, b = color
        for t in range(3):
            draw.rectangle([x1-t, y1-t, x2+t, y2+t], outline=(r, g, b))
        try:
            bb = draw.textbbox((x1, y1-22), lbl, font=font)
            tw, th = bb[2]-bb[0], bb[3]-bb[1]
        except Exception:
            tw, th = len(lbl)*8, 16
        draw.rectangle([x1, y1-th-4, x1+tw+4, y1], fill=(r, g, b))
        draw.text((x1+2, y1-th-2), lbl, fill=(255, 255, 255), font=font)
    return img_pil

def side_by_side(img_orig, dets_v11, dets_v5):
    w, h   = img_orig.size
    canvas = Image.new("RGB", (w*3+20, h+40), (18, 18, 18))
    draw   = ImageDraw.Draw(canvas)
    try:
        fh = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 20)
    except Exception:
        fh = ImageFont.load_default()
    for i, title in enumerate(["Original", "YOLOv11x", "YOLOv5x"]):
        draw.text((i*(w+10)+w//2-40, 8), title, fill=(200, 200, 200), font=fh)
    canvas.paste(img_orig.copy(),                                    (0,       40))
    canvas.paste(draw_boxes(img_orig.copy(), dets_v11, COLORS["YOLOv11x"], "v11"), (w+10,   40))
    canvas.paste(draw_boxes(img_orig.copy(), dets_v5,  COLORS["YOLOv5x"],  "v5"),  (2*w+20, 40))
    return canvas

# ── Main detect fn ────────────────────────────────────────────
def detect(image, model_choice, conf_thr, iou_thr):
    if image is None:
        return None, "⚠️ Upload an image first."
    img_np = np.array(image.convert("RGB"))
    t0 = time.time()
    dets_v11, dets_v5, lines = [], [], []

    if model_choice in ("YOLOv11x", "Both"):
        t1 = time.time(); dets_v11 = run_v11(img_np, conf_thr, iou_thr)
        lines.append(f"**YOLOv11x** — {len(dets_v11)} detection(s) · {(time.time()-t1)*1000:.0f} ms")

    if model_choice in ("YOLOv5x", "Both"):
        t1 = time.time(); dets_v5 = run_v5(img_np, conf_thr, iou_thr)
        lines.append(f"**YOLOv5x** — {len(dets_v5)} detection(s) · {(time.time()-t1)*1000:.0f} ms")

    lines.append(f"⏱ {(time.time()-t0)*1000:.0f} ms total · {DEVICE.upper()}")

    if model_choice == "YOLOv11x":
        out = draw_boxes(image.copy(), dets_v11, COLORS["YOLOv11x"], "v11")
    elif model_choice == "YOLOv5x":
        out = draw_boxes(image.copy(), dets_v5, COLORS["YOLOv5x"], "v5")
    else:
        out = side_by_side(image, dets_v11, dets_v5)

    return out, "\n\n".join(lines)

# ── Custom CSS ────────────────────────────────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Bebas+Neue&family=DM+Sans:wght@400;500&display=swap');

body, .gradio-container {
    background: #0f0f0f !important;
    color: #e8e8e8 !important;
    font-family: 'DM Sans', sans-serif !important;
}

#main-title {
    font-family: 'Bebas Neue', sans-serif !important;
    font-size: clamp(3rem, 8vw, 6rem) !important;
    letter-spacing: 0.06em;
    color: #ffffff;
    text-align: center;
    padding: 2rem 0 0.5rem 0;
    line-height: 1;
    border-bottom: 2px solid #f0a500;
    margin-bottom: 2rem;
}

/* Panels */
.gr-box, .gr-panel, .svelte-1gfkn6j, .block {
    background: #1a1a1a !important;
    border: 1px solid #2a2a2a !important;
    border-radius: 6px !important;
}

/* Radio buttons */
.gr-radio-row label {
    background: #222 !important;
    border: 1px solid #333 !important;
    border-radius: 4px !important;
    padding: 6px 14px !important;
    cursor: pointer;
    transition: all 0.15s;
}
.gr-radio-row label:hover { border-color: #f0a500 !important; }
.gr-radio-row input:checked + label,
.gr-radio-row label[data-checked] {
    border-color: #f0a500 !important;
    color: #f0a500 !important;
}

/* Sliders */
input[type=range] { accent-color: #f0a500; }

/* Detect button */
#detect-btn {
    background: #f0a500 !important;
    color: #0f0f0f !important;
    font-family: 'Bebas Neue', sans-serif !important;
    font-size: 1.3rem !important;
    letter-spacing: 0.1em;
    border: none !important;
    border-radius: 4px !important;
    height: 52px !important;
    transition: opacity 0.2s;
}
#detect-btn:hover { opacity: 0.85; }

/* Labels */
label span, .gr-label { color: #999 !important; font-size: 0.78rem !important; text-transform: uppercase; letter-spacing: 0.08em; }

/* Example thumbnails */
.gr-samples-table img { border-radius: 4px; border: 1px solid #333; }

/* Stats output */
#stats-out { font-size: 0.85rem !important; color: #aaa !important; }
"""

# ── UI ────────────────────────────────────────────────────────
with gr.Blocks(title="Mudflap Detector", css=CSS) as demo:

    gr.HTML('<div id="main-title">Mudflap Detector</div>')

    with gr.Row(equal_height=False):

        # ── Left column: inputs ───────────────────────────────
        with gr.Column(scale=1, min_width=300):
            inp_image = gr.Image(type="pil", label="Input Image", height=300)

            gr.HTML('<div style="font-size:0.72rem;color:#666;text-transform:uppercase;letter-spacing:.08em;margin:10px 0 6px">Example Images</div>')
            gr.Examples(
                examples=[
                    ["examples/example1.jpg", "Both", 0.60, 0.45],
                    ["examples/example2.jpg", "YOLOv11x", 0.60, 0.45],
                ],
                inputs=[inp_image,
                        gr.State("Both"),        # placeholder — wired below
                        gr.State(0.60),
                        gr.State(0.45)],
                label=None,
            )

            model_sel = gr.Radio(
                choices=["YOLOv11x", "YOLOv5x", "Both"],
                value="Both",
                label="Model",
            )

            with gr.Row():
                conf_slider = gr.Slider(0.10, 0.95, value=0.60, step=0.05, label="Confidence")
                iou_slider  = gr.Slider(0.10, 0.90, value=0.45, step=0.05, label="IoU")

            run_btn = gr.Button("DETECT", elem_id="detect-btn", variant="primary")

        # ── Right column: output ──────────────────────────────
        with gr.Column(scale=2, min_width=400):
            out_image = gr.Image(type="pil", label="Result", height=500)
            out_stats = gr.Markdown(elem_id="stats-out")

    run_btn.click(
        fn=detect,
        inputs=[inp_image, model_sel, conf_slider, iou_slider],
        outputs=[out_image, out_stats],
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)