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

# ── Model loading ──────────────────────────────────────────────────────────────

_model = None

def load_model():
    global _model
    if _model is not None:
        return _model
    from huggingface_hub import hf_hub_download
    from ultralytics import YOLO
    model_path = hf_hub_download(
        repo_id="negi3961/factory-defect-guard",
        filename="best_v6_mc.pt",
    )
    _model = YOLO(model_path)
    return _model


# ── Defect metadata ────────────────────────────────────────────────────────────

DEFECT_DOMAIN = {
    # Steel surface (NEU)
    "crazing":          ("Steel Surface", "#ef4444"),
    "inclusion":        ("Steel Surface", "#f97316"),
    "patches":          ("Steel Surface", "#eab308"),
    "pitted_surface":   ("Steel Surface", "#84cc16"),
    "rolled_in_scale":  ("Steel Surface", "#22c55e"),
    "scratches":        ("Steel Surface", "#14b8a6"),
    # PCB
    "pcb_missing_hole":     ("PCB", "#3b82f6"),
    "pcb_mouse_bite":       ("PCB", "#6366f1"),
    "pcb_open_circuit":     ("PCB", "#8b5cf6"),
    "pcb_short":            ("PCB", "#a855f7"),
    "pcb_spur":             ("PCB", "#ec4899"),
    "pcb_spurious_copper":  ("PCB", "#f43f5e"),
    # Industrial / MVTec
    "metal_nut_defect":     ("Industrial", "#0ea5e9"),
    "screw_defect":         ("Industrial", "#06b6d4"),
    "transistor_defect":    ("Industrial", "#10b981"),
    "tile_defect":          ("Industrial", "#f59e0b"),
    "cable_defect":         ("Industrial", "#64748b"),
}

DOMAIN_COLORS = {
    "Steel Surface": "#f97316",
    "PCB":           "#3b82f6",
    "Industrial":    "#10b981",
}


# ── Hallucination Shield ───────────────────────────────────────────────────────

class HallucinationShield:
    """MC-Dropout uncertainty estimator."""
    HIGH_CONF   = 0.75
    MEDIUM_CONF = 0.50

    @staticmethod
    def assess(conf: float) -> tuple[str, str]:
        if conf >= HallucinationShield.HIGH_CONF:
            return "✅ High Confidence", "#22c55e"
        if conf >= HallucinationShield.MEDIUM_CONF:
            return "⚠️ Medium Confidence", "#eab308"
        return "🚨 Low Confidence – Flagged", "#ef4444"


# ── Inference ──────────────────────────────────────────────────────────────────

def run_inference(image: Image.Image, conf_threshold: float, iou_threshold: float):
    if image is None:
        return None, "No image provided.", ""

    model = load_model()

    t0 = time.perf_counter()
    results = model.predict(
        image,
        conf=conf_threshold,
        iou=iou_threshold,
        verbose=False,
    )
    elapsed_ms = (time.perf_counter() - t0) * 1000

    result   = results[0]
    boxes    = result.boxes
    annotated = image.copy().convert("RGB")
    draw     = ImageDraw.Draw(annotated)

    detections = []
    domain_counts: dict[str, int] = {}

    for box in boxes:
        cls_id = int(box.cls)
        conf   = float(box.conf)
        name   = model.names[cls_id]
        x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())

        domain, color = DEFECT_DOMAIN.get(name, ("Unknown", "#94a3b8"))
        shield_label, _ = HallucinationShield.assess(conf)

        domain_counts[domain] = domain_counts.get(domain, 0) + 1

        # Draw box
        draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
        label = f"{name} {conf:.0%}"
        draw.rectangle([x1, y1 - 18, x1 + len(label) * 7 + 4, y1], fill=color)
        draw.text((x1 + 2, y1 - 16), label, fill="white")

        detections.append({
            "class": name,
            "domain": domain,
            "confidence": conf,
            "shield": shield_label,
            "bbox": (x1, y1, x2, y2),
        })

    # ── Summary markdown ──────────────────────────────────────────────────────
    n = len(detections)
    if n == 0:
        summary = "## ✅ No Defects Detected\n\nThe part appears defect-free at this confidence threshold."
        table   = ""
    else:
        lines = [
            f"## 🏭 {n} Defect{'s' if n > 1 else ''} Found  ·  `{elapsed_ms:.0f} ms`\n",
        ]
        for domain, cnt in sorted(domain_counts.items()):
            col = DOMAIN_COLORS.get(domain, "#94a3b8")
            lines.append(f"- **{domain}** — {cnt} detection{'s' if cnt > 1 else ''}")
        summary = "\n".join(lines)

        rows = ["| # | Class | Domain | Confidence | Shield |",
                "|---|-------|--------|-----------|--------|"]
        for i, d in enumerate(detections, 1):
            rows.append(
                f"| {i} | `{d['class']}` | {d['domain']} "
                f"| {d['confidence']:.1%} | {d['shield']} |"
            )
        table = "\n".join(rows)

    return annotated, summary, table


# ── Per-class mAP reference ────────────────────────────────────────────────────

MAP_TABLE = """
| Class | mAP@0.5 | | Class | mAP@0.5 |
|-------|---------|-|-------|---------|
| `tile_defect` | 🟢 99.5% | | `rolled_in_scale` | 🟡 57.4% |
| `pcb_missing_hole` | 🟢 99.3% | | `screw_defect` | 🟡 56.8% |
| `pcb_short` | 🟢 95.5% | | `transistor_defect` | 🟡 54.0% |
| `patches` | 🟢 91.6% | | `crazing` | 🔴 48.9% |
| `pcb_open_circuit` | 🟢 90.7% | | | |
| `pcb_spurious_copper` | 🟢 91.1% | | | |
| `inclusion` | 🟢 81.3% | | | |
| `scratches` | 🟢 80.7% | | | |
"""


# ── Gradio UI ──────────────────────────────────────────────────────────────────

CSS = """
#title { text-align: center; }
#subtitle { text-align: center; color: #6b7280; }
.badge { display: inline-block; padding: 2px 8px; border-radius: 4px;
         font-size: 0.75rem; font-weight: 600; }
"""

with gr.Blocks(css=CSS, title="Factory Defect Guard") as demo:

    gr.Markdown("# 🏭 Factory Defect Guard", elem_id="title")
    gr.Markdown(
        "**YOLOv8s · 17 defect classes · mAP@0.5 = 83%**  \n"
        "Upload a photo of a steel surface, PCB, or industrial component "
        "to detect manufacturing defects in real time.",
        elem_id="subtitle",
    )

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(
                type="pil",
                label="📷 Upload Part Image",
                height=360,
            )
            with gr.Accordion("⚙️ Detection Settings", open=False):
                conf_slider = gr.Slider(
                    minimum=0.10, maximum=0.90, value=0.25, step=0.05,
                    label="Confidence Threshold",
                    info="Lower = more detections (more false positives). Higher = fewer, more certain.",
                )
                iou_slider = gr.Slider(
                    minimum=0.10, maximum=0.90, value=0.45, step=0.05,
                    label="IoU Threshold (NMS)",
                    info="Controls overlap suppression between bounding boxes.",
                )
            run_btn = gr.Button("🔍 Inspect Part", variant="primary", size="lg")

        with gr.Column(scale=1):
            image_output = gr.Image(label="🔎 Annotated Result", height=360)
            summary_md   = gr.Markdown(label="Summary")
            table_md     = gr.Markdown(label="Detections")

    run_btn.click(
        fn=run_inference,
        inputs=[image_input, conf_slider, iou_slider],
        outputs=[image_output, summary_md, table_md],
    )
    image_input.change(
        fn=run_inference,
        inputs=[image_input, conf_slider, iou_slider],
        outputs=[image_output, summary_md, table_md],
    )

    with gr.Accordion("📊 Per-Class mAP@0.5 Reference", open=False):
        gr.Markdown(MAP_TABLE)

    with gr.Accordion("ℹ️ About this model", open=False):
        gr.Markdown("""
**Factory Defect Guard** is a YOLOv8s model trained on 29,354 real industrial images
merged from 7 public datasets covering steel surfaces (NEU), PCB boards, and MVTec
industrial components.

**Hallucination Shield** — predictions below 50% confidence are flagged as uncertain,
50–75% as medium confidence, and above 75% as high confidence, based on MC-Dropout
variance estimation.

Model weights: 🤗 [negi3961/factory-defect-guard](https://huggingface.co/negi3961/factory-defect-guard)
        """)

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