"""
STL Firearm/Knife Screener
---------------------------
Upload an STL file. The app rotates a virtual camera fully around the model
(covering all azimuths, plus dedicated top-down and bottom-up views), renders
each view as a 2D image, and runs the "Guns-100-11m" YOLOv11 detector
(wuhp/guns-100-11m) on every rendered frame. Results are aggregated into an
overall verdict plus an annotated gallery so a detection that only shows up
from one angle (e.g. the trigger guard only visible from the side) still
gets caught.

Run locally:
    pip install -r requirements.txt
    python app.py

Model weights are pulled automatically from the Hugging Face Hub the first
time the app runs (needs internet access) and cached locally afterwards.
"""

import os
import tempfile

import gradio as gr
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from matplotlib.colors import LightSource
import numpy as np
import pandas as pd
import trimesh
from PIL import Image, ImageDraw, ImageFont
from huggingface_hub import hf_hub_download
from ultralytics import YOLO

# --------------------------------------------------------------------------
# Model loading
# --------------------------------------------------------------------------

MODEL_REPO = "wuhp/guns-100-11m"
MODEL_FILE = "Guns-100-11m.pt"
CLASS_NAMES = {0: "Gun", 1: "Knife"}

_model = None


def get_model():
    """Lazily download (once) and cache the YOLO model."""
    global _model
    if _model is None:
        weights_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
        _model = YOLO(weights_path)
    return _model


# --------------------------------------------------------------------------
# STL -> multi-view rendering
# --------------------------------------------------------------------------

def build_view_angles(num_azimuths: int, num_elevations: int):
    """
    Build a list of (azimuth, elevation) pairs that fully cover the sphere
    around the object, including explicit top-down and bottom-up shots.
    """
    azimuths = np.linspace(0, 360, num_azimuths, endpoint=False)

    # Elevation ring between (but not including) the poles, e.g. for
    # num_elevations=3 -> [-45, 0, 45]
    if num_elevations > 0:
        elevations = np.linspace(-75, 75, num_elevations)
    else:
        elevations = np.array([0.0])

    views = []
    for elev in elevations:
        for az in azimuths:
            views.append((float(az), float(elev)))

    # Dedicated poles: straight down (top view) and straight up (bottom view)
    views.append((0.0, 90.0))   # top
    views.append((0.0, -90.0))  # bottom

    return views


def render_view(tris: np.ndarray, center: np.ndarray, scale: float,
                 azim: float, elev: float, img_size: int = 640) -> Image.Image:
    """Render a single (azim, elev) view of the mesh triangles to a PIL image."""
    fig = plt.figure(figsize=(img_size / 100, img_size / 100), dpi=100)
    ax = fig.add_subplot(111, projection="3d")
    ax.set_facecolor("white")

    ls = LightSource(azdeg=315, altdeg=55)
    facecolors = [(0.62, 0.62, 0.66, 1.0)] * len(tris)
    poly = Poly3DCollection(
        tris,
        facecolors=facecolors,
        edgecolor=(0.15, 0.15, 0.15, 0.4),
        linewidths=0.15,
        shade=True,
        lightsource=ls,
    )
    ax.add_collection3d(poly)

    ax.set_xlim(center[0] - scale, center[0] + scale)
    ax.set_ylim(center[1] - scale, center[1] + scale)
    ax.set_zlim(center[2] - scale, center[2] + scale)
    try:
        ax.set_box_aspect((1, 1, 1))
    except Exception:
        pass
    ax.set_axis_off()
    ax.view_init(elev=elev, azim=azim)
    fig.subplots_adjust(left=0, right=1, top=1, bottom=0)

    fig.canvas.draw()
    buf = np.asarray(fig.canvas.buffer_rgba())
    img = Image.fromarray(buf).convert("RGB")
    plt.close(fig)
    return img


def load_mesh_triangles(stl_path: str):
    mesh = trimesh.load(stl_path, force="mesh")
    if not isinstance(mesh, trimesh.Trimesh):
        # Scene with multiple geometries -> merge them
        mesh = trimesh.util.concatenate(mesh.dump())
    tris = mesh.vertices[mesh.faces]
    center = mesh.centroid
    scale = float(np.max(mesh.extents)) / 2.0 * 1.05  # small margin
    return tris, center, scale, mesh


# --------------------------------------------------------------------------
# Detection + aggregation
# --------------------------------------------------------------------------

def annotate_image(img: Image.Image, boxes, label_prefix=""):
    """Draw YOLO detection boxes on a copy of the image."""
    out = img.copy()
    draw = ImageDraw.Draw(out)
    try:
        font = ImageFont.load_default()
    except Exception:
        font = None

    for cls_id, conf, xyxy in boxes:
        x1, y1, x2, y2 = xyxy
        color = (220, 30, 30) if cls_id == 0 else (30, 120, 220)
        draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
        label = f"{CLASS_NAMES.get(cls_id, cls_id)} {conf:.2f}"
        draw.text((x1 + 3, max(0, y1 - 14)), label, fill=color, font=font)
    return out


def process_stl(stl_file, num_azimuths, num_elevations, conf_threshold, img_size,
                 progress=gr.Progress()):
    if stl_file is None:
        return None, "Please upload an STL file.", None

    progress(0, desc="Loading model...")
    model = get_model()

    progress(0.05, desc="Loading mesh...")
    tris, center, scale, mesh = load_mesh_triangles(stl_file)

    views = build_view_angles(int(num_azimuths), int(num_elevations))
    total = len(views)

    all_detections = []          # rows for the results table
    annotated_images = []        # (image, caption) for the gallery
    max_conf_per_class = {0: 0.0, 1: 0.0}
    any_hit = False

    for i, (az, el) in enumerate(views):
        progress((0.1 + 0.85 * i / total), desc=f"Rendering + detecting view {i+1}/{total}...")
        frame = render_view(tris, center, scale, az, el, img_size=int(img_size))

        result = model.predict(source=np.array(frame), conf=float(conf_threshold), verbose=False)[0]

        boxes = []
        for b in result.boxes:
            cls_id = int(b.cls.item())
            conf = float(b.conf.item())
            xyxy = [float(v) for v in b.xyxy[0].tolist()]
            boxes.append((cls_id, conf, xyxy))
            max_conf_per_class[cls_id] = max(max_conf_per_class[cls_id], conf)
            any_hit = True
            all_detections.append({
                "View": f"az={az:.0f} el={el:.0f}",
                "Class": CLASS_NAMES.get(cls_id, str(cls_id)),
                "Confidence": round(conf, 3),
            })

        if boxes:
            annotated = annotate_image(frame, boxes)
            caption = f"az={az:.0f}, el={el:.0f} | " + ", ".join(
                f"{CLASS_NAMES.get(c, c)} {p:.2f}" for c, p, _ in boxes
            )
            annotated_images.append((annotated, caption))

    progress(1.0, desc="Done")

    # Build verdict
    if any_hit:
        parts = []
        if max_conf_per_class[0] > 0:
            parts.append(f"GUN detected (max confidence {max_conf_per_class[0]:.2f})")
        if max_conf_per_class[1] > 0:
            parts.append(f"KNIFE detected (max confidence {max_conf_per_class[1]:.2f})")
        verdict = "\u26a0\ufe0f WEAPON-LIKE OBJECT DETECTED: " + " | ".join(parts)
        verdict += f"\n\nDetected in {len(annotated_images)} of {total} rendered views."
    else:
        verdict = f"\u2705 No firearm or knife detected across {total} rendered views."

    verdict += ("\n\nNote: this is a heuristic screen based on rendered silhouettes "
                "of the mesh, not a photo of a real object. Treat results as advisory, "
                "not a certified determination.")

    if not annotated_images:
        # still show a couple of representative views so the user sees *something*
        sample_idxs = [0, total // 2] if total > 1 else [0]
        for idx in sample_idxs:
            az, el = views[idx]
            frame = render_view(tris, center, scale, az, el, img_size=int(img_size))
            annotated_images.append((frame, f"az={az:.0f}, el={el:.0f} (no detection)"))

    df = pd.DataFrame(all_detections) if all_detections else pd.DataFrame(
        columns=["View", "Class", "Confidence"])

    return annotated_images, verdict, df


# --------------------------------------------------------------------------
# Gradio UI
# --------------------------------------------------------------------------

with gr.Blocks(title="STL Firearm/Knife Screener") as demo:
    gr.Markdown(
        "# 🔫🔪 STL Firearm / Knife Screener\n"
        "Upload a 3D model (`.stl`). The app spins a virtual camera all the way "
        "around it — including straight-down and straight-up shots — renders each "
        "view, and runs the **Guns-100-11m** YOLOv11 detector on every frame."
    )

    with gr.Row():
        with gr.Column(scale=1):
            stl_input = gr.File(label="Upload STL file", file_types=[".stl"], type="filepath")
            num_azimuths = gr.Slider(4, 16, value=8, step=1, label="Azimuth steps per elevation ring (more = finer 360° coverage, slower)")
            num_elevations = gr.Slider(1, 5, value=3, step=1, label="Elevation rings between the poles")
            conf_threshold = gr.Slider(0.05, 0.95, value=0.25, step=0.05, label="Detection confidence threshold")
            img_size = gr.Slider(320, 960, value=640, step=32, label="Render resolution (px)")
            run_btn = gr.Button("Run 360° Scan", variant="primary")

        with gr.Column(scale=2):
            verdict_out = gr.Textbox(label="Verdict", lines=4)
            gallery_out = gr.Gallery(label="Views with detections", columns=3, height=500)
            table_out = gr.Dataframe(label="All detections", wrap=True)

    run_btn.click(
        fn=process_stl,
        inputs=[stl_input, num_azimuths, num_elevations, conf_threshold, img_size],
        outputs=[gallery_out, verdict_out, table_out],
    )

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