# app.py — Hugging Face Spaces demo (Gradio) over the shipped detector.
# Same model, preprocessing and evidence-map rendering as `predict.py --serve`;
# this file only adds the hosted UI. Run locally: `python app.py`.
import json
import math
import os
import sys
import time

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "src"))

import gradio as gr  # noqa: E402
import torch  # noqa: E402

try:                      # Hugging Face ZeroGPU: the GPU is attached only inside
    import spaces         # @spaces.GPU-decorated calls; a no-op elsewhere.
except ImportError:       # local / plain-CPU runs
    class spaces:  # noqa: N801
        @staticmethod
        def GPU(fn=None, **_):  # noqa: N802
            return fn if fn is not None else (lambda f: f)

from inference.predict import _resolve_device  # noqa: E402
from inference.server import render_overlay, score_image  # noqa: E402
from train.checkpoint import load_ensemble_checkpoint  # noqa: E402

ROOT = os.path.dirname(os.path.abspath(__file__))
WEIGHTS = os.environ.get("WEIGHTS", os.path.join(ROOT, "weights.pt"))
# On ZeroGPU torch.cuda.is_available() is True at import time even though no GPU is
# attached yet; .to("cuda") here is what the spaces runtime intercepts and replays.
# INFER_DEVICE=cpu forces CPU inference (no GPU quota consumed, no ZeroGPU queue).
DEVICE = _resolve_device(os.environ.get("INFER_DEVICE", "auto"))
GPU_SECONDS = 8 if DEVICE == "cuda" else 0
MEMBERS = load_ensemble_checkpoint(WEIGHTS, map_location="cpu")
for m, _ in MEMBERS:
    m.to(DEVICE).eval()
torch.set_num_threads(max(1, os.cpu_count() or 1))

SAMPLE_DIR = os.path.join(ROOT, "samples")
with open(os.path.join(SAMPLE_DIR, "samples.json")) as f:
    SAMPLES = json.load(f)                      # see samples/ATTRIBUTION.md
SAMPLE_PATHS = [os.path.join(SAMPLE_DIR, s["file"]) for s in SAMPLES]

# Display-only operating point. 0.62 is the balanced-error point (FPR == FNR, ~2%) of
# the ensemble on the clean holdout and the lowest total error on every corruption
# family; 0.5 over-flags real photos ~2:1. Neither the CLI nor the JSON output moves.
VERDICT_THRESHOLD = 0.62
_THR_LOGIT = math.log(VERDICT_THRESHOLD / (1.0 - VERDICT_THRESHOLD))


def decision_score(p):
    """Recentre the raw probability so the verdict boundary sits at 0.5: a shift in
    logit space (identical to moving the classifier's bias term). Monotone, kink-free,
    0 -> 0, VERDICT_THRESHOLD -> 0.5, 1 -> 1; the extremes barely move."""
    p = min(max(p, 1e-6), 1.0 - 1e-6)
    return 1.0 / (1.0 + math.exp(-(math.log(p / (1.0 - p)) - _THR_LOGIT)))


def _score(img):
    # Timed INSIDE the GPU-scoped call, so on ZeroGPU this is the model's own time
    # (preprocess + two forwards; score_image's float()/.cpu() force the CUDA sync),
    # excluding queue wait and the runtime's replay of the .to("cuda") moves.
    t0 = time.perf_counter()
    entry, evid = score_image(MEMBERS, img, DEVICE, explain=True)
    entry["infer_s"] = time.perf_counter() - t0
    return entry, evid


if GPU_SECONDS:   # only the forward pass holds the GPU; overlay rendering is CPU work
    _score = spaces.GPU(duration=GPU_SECONDS)(_score)


def detect(img):
    if img is None:
        return "", None, None, ""
    t0 = time.perf_counter()
    entry, evid = _score(img)
    dt = time.perf_counter() - t0
    p = entry["pred"]
    s = decision_score(p)
    overlay = render_overlay(img, evid, max_side=640)
    if s >= 0.5:
        verdict = f"## 🤖 AI-generated  ·  score {s:.2f}"
    else:
        verdict = f"## 📷 Authentic  ·  score {s:.2f}"
    # explain_prediction's rationale opens with its own 0.5-threshold verdict; keep only
    # the "the detector's evidence is ..." clause so the two verdicts can't disagree.
    where = entry["rationale"].split("; ", 1)[-1]
    return (verdict, {"AI-generated": s, "authentic": 1.0 - s}, overlay,
            f"{where[0].upper()}{where[1:]} Raw ensemble p(AI) = {p:.3f}; the score "
            f"recentres the {VERDICT_THRESHOLD} operating point to 0.5. "
            f"[{DEVICE}: model {entry['infer_s']:.2f}s, {dt:.2f}s incl. queue]")


def pick_sample(evt: gr.SelectData):
    return SAMPLE_PATHS[evt.index]      # path, not PIL: no lossy re-encode in transit


DESCRIPTION = """
Upload, drag-drop or paste an image — or pick one of the samples in the sidebar. The
detector returns the probability the image is AI-generated and a signed evidence map
(**red** = fake evidence, **blue** = real evidence, grayscale = neutral). The score is
the ensemble probability recentred so that 0.5 is the decision boundary: the raw model
balances its two error rates at p = 0.62 on the held-out set, and the score shifts that
point to 0.5 (a logit offset — the ranking of images is unchanged; the raw p is in the
rationale).

Frozen DINOv2-B/14 read at blocks 8 and 10, two small transformer heads, logits
averaged — the exact weights from the submission repo. Source, weights and full evaluation:
[github.com/samt41/techjam_aigsee](https://github.com/samt41/techjam_aigsee).
"""

with gr.Blocks(title="AIGC image detector — TechJam 2026 Track 5") as demo:
    with gr.Sidebar(label="Samples", width=300):
        gr.Markdown("### Try a sample\nClick a thumbnail to load it. Half are AI-generated, "
                    "half are real photographs — labels are revealed in the caption.")
        gallery = gr.Gallery(value=[(p, s["label"]) for p, s in zip(SAMPLE_PATHS, SAMPLES)],
                             columns=2, height="auto", object_fit="cover",
                             allow_preview=False, show_label=False)
        gr.Markdown("AI samples: Community Forensics (CC-BY-4.0). Real photos: Wikimedia "
                    "Commons (CC0). Full credits in `samples/ATTRIBUTION.md`.")
    gr.Markdown("# AIGC image detector — TechJam 2026 Track 5")
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column():
            # format="png": gradio's default is lossy WebP, which would re-encode
            # every upload before the model sees it.
            inp = gr.Image(type="pil", format="png",
                           label="Image (upload, drag-drop or paste)",
                           sources=["upload", "clipboard"])
            btn = gr.Button("Detect", variant="primary")
        with gr.Column():
            ver = gr.Markdown("")
            lab = gr.Label(label="Score (0.5 = decision boundary)", num_top_classes=2)
            ovl = gr.Image(label="Evidence map (red = fake, blue = real)")
            rat = gr.Textbox(label="Rationale", lines=2)
    outputs = [ver, lab, ovl, rat]
    btn.click(detect, inputs=inp, outputs=outputs)
    inp.change(detect, inputs=inp, outputs=outputs)
    gallery.select(pick_sample, inputs=None, outputs=inp)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))