import json
from pathlib import Path
import spaces
import numpy as np
import torch
import torchvision.transforms as T
import gradio as gr
from PIL import Image
from matplotlib import colormaps

from model import SceneCNN, grad_cam

ROOT = Path(__file__).parent

# ---------- load model ----------
meta = json.loads((ROOT / "cnn_intel_meta.json").read_text())
CLASSES  = meta["class_names"]
IMG_SIZE = meta["img_size"]

model = SceneCNN(base=meta["base"], num_classes=len(CLASSES))
model.load_state_dict(torch.load(ROOT / "cnn_intel.pth", map_location="cpu"))
model.eval()

preprocess = T.Compose([
    T.Resize((IMG_SIZE, IMG_SIZE)),
    T.ToTensor(),
    T.Normalize(meta["mean"], meta["std"]),
])

@spaces.GPU
# ---------- prediction ----------
def predict(img):
    if img is None:
        return None, None
    img = img.convert("RGB")
    cam, probs = grad_cam(model, preprocess(img))

    show = 384
    base = np.asarray(img.resize((show, show))) / 255.0
    cam_big = np.asarray(Image.fromarray((cam * 255).astype(np.uint8)).resize((show, show), Image.BILINEAR)) / 255.0
    heat = colormaps["jet"](cam_big)[..., :3]
    overlay = (0.55 * base + 0.45 * heat).clip(0, 1)

    label = {c: float(p) for c, p in zip(CLASSES, probs)}
    return label, (overlay * 255).astype(np.uint8)

# ---------- metrics ----------
def metrics_markdown():
    p = ROOT / "metrics.json"
    if not p.exists():
        return "_metrics.json not found._"
    m = json.loads(p.read_text())
    cm = np.array(m["confusion_matrix"])
    tp = cm.diagonal()
    precision = tp / cm.sum(axis=0).clip(min=1)
    recall    = tp / cm.sum(axis=1).clip(min=1)
    f1        = 2 * precision * recall / (precision + recall + 1e-12)

    head = [f"**Test accuracy:** {m['test_accuracy']:.2%}",
            f"**Macro F1:** {m['macro_f1']:.3f}"]
    if m.get("best_val_accuracy"):
        head.append(f"**Best val accuracy:** {m['best_val_accuracy']:.2%}")
    if m.get("params"):
        head.append(f"**Parameters:** {m['params']:,}")

    out = ["  ·  ".join(head), "",
           "| Class | Precision | Recall | F1 | Test images |",
           "|---|---|---|---|---|"]
    for i, c in enumerate(CLASSES):
        out.append(f"| {c} | {precision[i]:.3f} | {recall[i]:.3f} | {f1[i]:.3f} | {cm[i].sum()} |")

    # top confusions
    off = [(cm[i, j], CLASSES[i], CLASSES[j]) for i in range(len(CLASSES))
           for j in range(len(CLASSES)) if i != j]
    out += ["", "**Most confused (actual → predicted):** " +
            ", ".join(f"{a} → {b} ({n})" for n, a, b in sorted(off, reverse=True)[:4])]
    return "\n".join(out)

PLOTS = [
    ("10_curves.png",               "Training curves"),
    ("11_confusion_matrix.png",     "Confusion matrix"),
    ("12_per_class_accuracy.png",   "Per-class accuracy"),
    ("13_confidence.png",           "Confidence: correct vs wrong"),
    ("14_mistakes.png",             "Most confident mistakes"),
    ("15_gradcam.png",              "Grad-CAM per class"),
    ("16_gradcam_mistakes.png",     "Grad-CAM on mistakes"),
    ("08_feature_maps_after.png",   "Feature maps after training"),
    ("09_filters_after.png",        "First-layer filters after training"),
    ("04_augmentation.png",         "Augmentation"),
    ("01_samples.png",              "Dataset samples"),
    ("02_class_distribution.png",   "Class distribution"),
]
gallery_items = [(str(ROOT / "images" / f), cap) for f, cap in PLOTS if (ROOT / "images" / f).exists()]
examples = [[str(p)] for p in sorted((ROOT / "examples").glob("*.jpg"))] if (ROOT / "examples").exists() else []

# ---------- UI ----------
with gr.Blocks(title="Intel Scene CNN", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        "# Intel Scene Classifier: a CNN built from scratch\n"
        "Upload a photo of a **building, forest, glacier, mountain, sea or street**. "
        "The CNN predicts the scene, and Grad-CAM shows which pixels drove the decision. "
        "No pretrained backbone: 4 conv blocks trained from scratch on ~14k images."
    )
    with gr.Row():
        inp = gr.Image(type="pil", label="Upload a scene", height=384)
        with gr.Column():
            out_label = gr.Label(num_top_classes=6, label="Prediction")
            out_cam = gr.Image(label="Grad-CAM (red = most important)", height=384)
    inp.change(predict, inputs=inp, outputs=[out_label, out_cam])
    if examples:
        gr.Examples(examples=examples, inputs=inp, cache_examples=False)

    gr.Markdown("> Note: the model always picks one of the 6 classes, even for a photo of a cat.")

    gr.Markdown("## Test-set metrics (seg_test, 3,000 held-out images)")
    gr.Markdown(metrics_markdown())

    if gallery_items:
        gr.Markdown("## Training and evaluation plots")
        gr.Gallery(value=gallery_items, columns=2, height="auto", object_fit="contain")

    gr.Markdown(
        "## Architecture\n"
        "`[Conv3x3-BN-ReLU] x2 → MaxPool → Dropout2d` × 4 blocks "
        f"(channels {meta['base']}→{meta['base']*2}→{meta['base']*4}→{meta['base']*8}) "
        "→ Global Avg Pool → Dropout → FC 128 → ReLU → FC 6"
    )

demo.launch(show_api=False, ssr_mode=False)