import gradio as gr
import torch
import numpy as np
from PIL import Image
from ultralytics import YOLO
import cv2

# ─── Load Model ───────────────────────────────────────────────────────────────
MODEL_PATH  = "yolov8s_best.pt"
CLASS_NAMES = [
    "D00 — Longitudinal Crack",
    "D10 — Transverse Crack",
    "D20 — Alligator Crack",
    "D40 — Pothole",
]
CLASS_COLORS = {
    "D00 — Longitudinal Crack": "#3498db",
    "D10 — Transverse Crack":   "#2ecc71",
    "D20 — Alligator Crack":    "#e67e22",
    "D40 — Pothole":            "#e74c3c",
}
COUNTRY_POTHOLE = {
    "🇮🇳 India":          "⬛⬛⬛⬛⬛ Very High",
    "🇳🇴 Norway":         "⬛⬛⬛⬛  High",
    "🇺🇸 United States":  "⬛⬛⬛   Moderate",
    "🇨🇿 Czech Republic": "⬛⬛⬛   Moderate",
    "🇨🇳 China":          "⬛⬛    Low-Moderate",
    "🇯🇵 Japan":          "⬛     Low",
}

model = YOLO(MODEL_PATH)

# ─── Inference Function ───────────────────────────────────────────────────────
def detect(image, conf_thresh, iou_thresh, country):
    img_np = np.array(image)

    results = model.predict(
        source    = img_np,
        conf      = conf_thresh,
        iou       = iou_thresh,
        imgsz     = 640,
        device    = 0 if torch.cuda.is_available() else "cpu",
        verbose   = False,
    )[0]

    # ── Annotated image ───────────────────────────────────────────────────────
    annotated = results.plot()
    annotated_rgb = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)

    # ── Count detections ──────────────────────────────────────────────────────
    class_counts = {name: 0 for name in CLASS_NAMES}
    total        = len(results.boxes)
    pothole_count = 0

    for box in results.boxes:
        cls_id = int(box.cls.item())
        class_counts[CLASS_NAMES[cls_id]] += 1
        if cls_id == 3:
            pothole_count += 1

    # ── Detection table ───────────────────────────────────────────────────────
    det_table = "### 📊 Detection Summary\n\n"
    det_table += "| Class | Count | % of Total |\n|---|---|---|\n"
    for name, count in class_counts.items():
        pct = count / max(total, 1) * 100
        color_dot = "🔵" if "D00" in name else "🟢" if "D10" in name else "🟠" if "D20" in name else "🔴"
        det_table += f"| {color_dot} {name} | {count} | {pct:.1f}% |\n"
    det_table += f"\n**Total Detections: {total}**"

    # ── Pothole report box ────────────────────────────────────────────────────
    pct_pothole = pothole_count / max(total, 1) * 100
    if pct_pothole > 30:
        severity = "🔴 CRITICAL — Immediate repair needed"
    elif pct_pothole > 10:
        severity = "🟠 MODERATE — Schedule maintenance"
    elif pothole_count > 0:
        severity = "🟡 LOW — Monitor road surface"
    else:
        severity = "🟢 NONE DETECTED — Road surface OK"

    pothole_box = f"""### 🕳️ Pothole (D40) Report

| Field | Value |
|---|---|
| Pothole Count | **{pothole_count}** |
| % of Detections | **{pct_pothole:.1f}%** |
| Severity | {severity} |
| Avg Confidence | {"N/A" if pothole_count == 0 else f"{np.mean([b.conf.item() for b in results.boxes if int(b.cls.item())==3]):.2f}"} |
"""

    # ── Country pothole context ───────────────────────────────────────────────
    country_info = COUNTRY_POTHOLE.get(country, "—")
    country_box = f"""### 🌍 Country Context — {country}

| Field | Value |
|---|---|
| Selected Country | {country} |
| Pothole Density (RDD2022) | {country_info} |
| Dominant Damage Types | {"D00 / D10" if "Japan" in country else "D40 / D20" if "India" in country else "D10 / D40"} |
| Collection Method | {"Smartphone" if "India" in country or "Japan" in country or "Czech" in country else "Street View" if "United" in country else "Camera / Drone"} |
"""

    return Image.fromarray(annotated_rgb), det_table, pothole_box, country_box


# ─── Gradio UI ────────────────────────────────────────────────────────────────
with gr.Blocks(
    title="🛣️ Road Damage Detector — RDD2022",
    theme=gr.themes.Soft(primary_hue="blue"),
    css="""
    .title-box { text-align: center; padding: 20px; }
    .result-box { border-radius: 10px; padding: 10px; }
    footer { display: none !important; }
    """
) as demo:

    # Header
    gr.HTML("""
    <div class="title-box">
        <h1>🛣️ Road Damage Detection</h1>
        <p style="color:gray;">YOLOv8s · Trained on RDD2022 · 4 Damage Classes · 6 Countries</p>
        <p>
            <span style="background:#3498db;color:white;padding:3px 8px;border-radius:4px;margin:2px;">🔵 D00 Longitudinal</span>
            <span style="background:#2ecc71;color:white;padding:3px 8px;border-radius:4px;margin:2px;">🟢 D10 Transverse</span>
            <span style="background:#e67e22;color:white;padding:3px 8px;border-radius:4px;margin:2px;">🟠 D20 Alligator</span>
            <span style="background:#e74c3c;color:white;padding:3px 8px;border-radius:4px;margin:2px;">🔴 D40 Pothole</span>
        </p>
    </div>
    """)

    with gr.Row():
        # ── Left column: inputs ───────────────────────────────────────────────
        with gr.Column(scale=1):
            input_image = gr.Image(
                type="pil",
                label="📷 Upload Road Image",
                height=350,
            )
            with gr.Row():
                conf_slider = gr.Slider(
                    minimum=0.1, maximum=0.9, value=0.25, step=0.05,
                    label="Confidence Threshold"
                )
                iou_slider = gr.Slider(
                    minimum=0.1, maximum=0.9, value=0.45, step=0.05,
                    label="IoU Threshold"
                )
            country_dropdown = gr.Dropdown(
                choices=list(COUNTRY_POTHOLE.keys()),
                value="🇮🇳 India",
                label="🌍 Country (photo origin — for pothole context)",
            )
            run_btn = gr.Button("🔍 Detect Damage", variant="primary", size="lg")

        # ── Right column: outputs ─────────────────────────────────────────────
        with gr.Column(scale=1):
            output_image = gr.Image(
                label="📍 Detection Result",
                height=350,
            )

    with gr.Row():
        with gr.Column():
            det_output     = gr.Markdown(label="Detection Summary")
        with gr.Column():
            pothole_output = gr.Markdown(label="Pothole Report")
        with gr.Column():
            country_output = gr.Markdown(label="Country Context")

    # ── Examples ──────────────────────────────────────────────────────────────
    gr.Examples(
        examples=[
            ["examples/india_test_image.jpg",  0.25, 0.45, "🇮🇳 India"],
            ["examples/China_Drone_000253.jpg",  0.25, 0.45, "🇯🇵 Japan"],
            ["examples/United_States_004798.jpg", 0.25, 0.45, "🇺🇸 United States"],
            ["examples/Czech_test_image.jpg", 0.25, 0.45, "🇨🇿 Czech Republic"],
            ["examples/China_Drone_000295.jpg", 0.25, 0.45, "🇨🇳 China"],
            ["examples/norway_road_test.jpg", 0.25, 0.45, "🇳🇴 Norway"]
        ],
        inputs=[input_image, conf_slider, iou_slider, country_dropdown],
        label="📂 Example Images",
    )

    # ── Footer ────────────────────────────────────────────────────────────────
    gr.HTML("""
    <div style="text-align:center; margin-top:20px; color:gray; font-size:13px;">
        Model: YOLOv8s · Dataset: RDD2022 (47,420 images · 6 countries) ·
        Classes: D00 / D10 / D20 / D40 ·
        <a href="https://arxiv.org/abs/2209.08538" target="_blank">Paper</a>
    </div>
    """)

    # ── Wire up ───────────────────────────────────────────────────────────────
    run_btn.click(
        fn      = detect,
        inputs  = [input_image, conf_slider, iou_slider, country_dropdown],
        outputs = [output_image, det_output, pothole_output, country_output],
    )
    input_image.change(
        fn      = detect,
        inputs  = [input_image, conf_slider, iou_slider, country_dropdown],
        outputs = [output_image, det_output, pothole_output, country_output],
    )

if __name__ == "__main__":
    demo.launch(
        server_name = "0.0.0.0",
        server_port = 7860,
        share       = True,      # generates public gradio.live link
    )