"""Hugging Face Space: upload a PCB image, get defect boxes from the INT8 ONNX model.

Pure ONNX Runtime + OpenCV (no PyTorch / Ultralytics), so the Space is small and
runs on any hardware tier's CPU. Put the quantized model next to this file as
model_int8.onnx (or set MODEL_PATH). Optional: a few test-split images in examples/.
"""
import ast
import os
import time
from pathlib import Path

import cv2
import gradio as gr
import numpy as np
import onnxruntime as ort

try:  # Hugging Face ZeroGPU hardware expects at least one @spaces.GPU function to exist.
    import spaces

    @spaces.GPU
    def _unused_gpu_placeholder():
        """Never called: inference below runs on CPU only."""
        return None
except ImportError:  # local runs / other hardware
    pass

HERE = Path(__file__).parent
MODEL_PATH = os.environ.get("MODEL_PATH", str(HERE / "model_int8.onnx"))
IMGSZ = int(os.environ.get("IMGSZ", "1280"))
NMS_IOU = 0.7  # Ultralytics default
FALLBACK_NAMES = ["missing_hole", "mouse_bite", "open_circuit", "short", "spur", "spurious_copper"]
COLORS = [(56, 56, 255), (151, 157, 255), (31, 112, 255), (29, 178, 255), (49, 210, 207), (10, 249, 72)]  # BGR

opts = ort.SessionOptions()
opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
session = ort.InferenceSession(MODEL_PATH, opts, providers=["CPUExecutionProvider"])
INPUT_NAME = session.get_inputs()[0].name
_meta = session.get_modelmeta().custom_metadata_map
NAMES = ast.literal_eval(_meta["names"]) if "names" in _meta else dict(enumerate(FALLBACK_NAMES))


def letterbox(img, size):
    """Resize keeping aspect ratio and pad to size x size (gray 114), like Ultralytics."""
    h, w = img.shape[:2]
    r = min(size / h, size / w)
    nh, nw = round(h * r), round(w * r)
    top, left = (size - nh) // 2, (size - nw) // 2
    out = np.full((size, size, 3), 114, dtype=np.uint8)
    out[top:top + nh, left:left + nw] = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_LINEAR)
    return out, r, left, top


def predict(bgr, conf):
    """Return a list of (class_id, score, x1, y1, x2, y2) in original-image pixels."""
    padded, r, dx, dy = letterbox(bgr, IMGSZ)
    x = np.ascontiguousarray(padded[:, :, ::-1].transpose(2, 0, 1), dtype=np.float32)[None] / 255.0
    out = session.run(None, {INPUT_NAME: x})[0][0].T          # (anchors, 4 + num_classes)
    scores = out[:, 4:]
    cls = scores.argmax(1)
    score = scores[np.arange(len(cls)), cls]
    keep = score >= conf
    if not keep.any():
        return []
    xywh, cls, score = out[keep, :4], cls[keep], score[keep]
    h, w = bgr.shape[:2]
    x1 = np.clip((xywh[:, 0] - xywh[:, 2] / 2 - dx) / r, 0, w)
    y1 = np.clip((xywh[:, 1] - xywh[:, 3] / 2 - dy) / r, 0, h)
    x2 = np.clip((xywh[:, 0] + xywh[:, 2] / 2 - dx) / r, 0, w)
    y2 = np.clip((xywh[:, 1] + xywh[:, 3] / 2 - dy) / r, 0, h)
    boxes = np.stack([x1, y1, x2 - x1, y2 - y1], 1)
    idx = cv2.dnn.NMSBoxesBatched(boxes.tolist(), score.tolist(), cls.tolist(), conf, NMS_IOU)
    idx = np.asarray(idx, dtype=int).reshape(-1)
    return [(int(cls[i]), float(score[i]), float(x1[i]), float(y1[i]), float(x2[i]), float(y2[i])) for i in idx]


def draw(bgr, dets):
    out = bgr.copy()
    thick = max(2, round(max(out.shape[:2]) / 600))
    for c, s, x1, y1, x2, y2 in dets:
        color = COLORS[c % len(COLORS)]
        cv2.rectangle(out, (int(x1), int(y1)), (int(x2), int(y2)), color, thick)
        cv2.putText(out, f"{NAMES[c]} {s:.2f}", (int(x1), max(20, int(y1) - 8)),
                    cv2.FONT_HERSHEY_SIMPLEX, thick / 2.5, color, thick, cv2.LINE_AA)
    return out


def detect(image, conf):
    if image is None:
        return None, [], ""
    bgr = np.ascontiguousarray(image[:, :, ::-1])
    t0 = time.perf_counter()
    dets = predict(bgr, conf)
    ms = (time.perf_counter() - t0) * 1000
    rows = [[NAMES[c], round(s, 3), int(x1), int(y1), int(x2), int(y2)] for c, s, x1, y1, x2, y2 in dets]
    status = f"{len(rows)} defects found in {ms:.0f} ms (CPU, ONNX Runtime, preprocessing + model + NMS)"
    return draw(bgr, dets)[:, :, ::-1], rows, status


# Example images: an examples/ folder, or .jpg files uploaded next to app.py.
EXAMPLES = {p.stem: str(p) for p in sorted([*HERE.glob("examples/*.jpg"), *HERE.glob("*.jpg")])}


def run_example(name, conf):
    """Load a bundled test image and detect on it in one step."""
    if not name:
        return None, None, [], ""
    rgb = np.ascontiguousarray(cv2.imread(EXAMPLES[name])[:, :, ::-1])
    annotated, rows, status = detect(rgb, conf)
    return rgb, annotated, rows, status

with gr.Blocks(title="PCB defect detection (INT8, CPU)") as demo:
    gr.Markdown(
        "# PCB defect detection\n"
        "YOLO11n quantized to INT8 with ONNX Runtime (detection head kept in FP32), running on CPU. "
        "Trained on PKU-Market-PCB, a public research dataset, not production factory data. "
        "Defects are tiny: zoom into the result to see the boxes."
    )
    with gr.Row():
        inp = gr.Image(type="numpy", label="PCB image", format="jpeg")
        out = gr.Image(type="numpy", label="Detections", format="jpeg")
    conf = gr.Slider(0.05, 0.9, value=0.25, step=0.05, label="Confidence threshold")
    status = gr.Markdown()
    table = gr.Dataframe(headers=["class", "conf", "x1", "y1", "x2", "y2"], label="Boxes")
    if EXAMPLES:
        pick = gr.Radio(list(EXAMPLES), value=None, label="Or pick a test-split example, then press Detect")
    gr.Button("Detect", variant="primary").click(detect, [inp, conf], [out, table, status])
    if EXAMPLES:
        pick.change(run_example, [pick, conf], [inp, out, table, status])

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