"""Hugging Face Space demo for Problem #10 (Explainable & Auditable Alloy Defect
Detection). Reduced-scope deliverable: upload a micrograph -> defect prediction +
Grad-CAM heatmap, computed live. Conformal Risk Control results are shown as the
pre-computed chart/numbers from the notebook rather than a live-recalibrated model,
per the scope committed in wiki/synthesis/problem10-explainable-auditable-ai.md."""
import json
import os

import gradio as gr
import torch

from model import CLASS_NAMES, GradCAM, eval_transform, load_model, overlay_heatmap

DEVICE = "cpu"
CHECKPOINT_PATH = "model.pt"

model = load_model(CHECKPOINT_PATH, device=DEVICE)
target_layer = model.features[-1]
gradcam = GradCAM(model, target_layer)


def predict(image):
    if image is None:
        return None, {}
    tensor = eval_transform(image.convert("RGB")).unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        logits = model(tensor)
        probs = logits.sigmoid()[0]
    top_class = int(probs.argmax())
    cam, _ = gradcam(tensor, torch.tensor([top_class]))
    overlay = overlay_heatmap(tensor[0], cam[0])
    label_scores = {CLASS_NAMES[i]: float(probs[i]) for i in range(len(CLASS_NAMES))}
    return overlay, label_scores


SPURIOUS_SETUP = """
    ### We planted a shortcut. The obvious test missed it.

    Grad-CAM tells you *where* the model looked, not whether what it found was real. So
    we planted a fake: a second model was trained from scratch on data where every image
    of one defect class (`scratches`) carries a crude synthetic "batch ID stamp" in the
    corner. Nothing about that stamp is metallurgical — it is exactly the kind of
    acquisition artifact (line ID, scale bar, microscope watermark) Problem #10 warns
    about.

    *The strategy is the course's:* this is **exercise 5 of the course Grad-CAM
    notebook** — "create a challenge set and test for shortcut learning" —
    proposed there but left unbuilt. Building it out is our contribution.

    Below: four images that are **not** `scratches`, with the stamp painted on at test
    time. **Top row** is what the model sees; **bottom row** is where it looked. (Each
    row's captions sit *above* its own images, so the `Grad-CAM` labels belong to the
    row beneath them.) The `p=` values are the model's `scratches` probability for
    images containing no scratch defect at all. The stamp pulls attention in every
    panel — which is what the third test below quantifies.
    """


def _spurious_markdown():
    """Build the shortcut-learning results from the audit's own saved numbers, so the
    text can never drift out of sync with the experiment that produced it."""
    try:
        with open("assets/spurious-results.json") as fh:
            r = json.load(fh)
    except (OSError, ValueError):
        return "_Results file unavailable._"
    ratio = r["gradcam_mass_in_stamp_corner"] / r["stamp_corner_area_fraction"]
    return f"""
        | Test | Result | |
        |---|---|---|
        | `scratches` recall, stamp present → removed | {r['recall_poisoned_class_stamped']:.0%} → **{r['recall_poisoned_class_clean']:.0%}** | no collapse |
        | Other defects called `scratches` once stamped | {r['baseline_other_classes_called_poisoned']:.0%} → **{r['hijack_rate_other_classes_when_stamped']:.1%}** | hijacked |
        | Grad-CAM mass inside the stamp corner | **{r['gradcam_mass_in_stamp_corner']:.1%}** on {r['stamp_corner_area_fraction']:.1%} of pixels | **{ratio:.0f}×** |

        **Strip the stamp off and recall stays at {r['recall_poisoned_class_clean']:.0%}.**
        The model never became *dependent* on the artifact — it learned the real defect
        signal and the stamp in parallel. That is what makes this dangerous: the test almost
        anyone would run first (remove the suspicious variable, check whether performance
        drops) returns a completely clean bill of health.

        What the artifact did do is become an *independently sufficient* cue. Painting it
        onto a different defect flips the model to `scratches`
        **{r['hijack_rate_other_classes_when_stamped']:.1%}** of the time against a
        **{r['baseline_other_classes_called_poisoned']:.0%}** baseline on identical unstamped
        images — and it captures **{ratio:.0f}×** more Grad-CAM attention than its
        area warrants, firing on the stamp even in cases where the prediction did *not* flip.

        So the honest statement to a plant manager isn't "the model is broken." It's: *this
        model attends to something with no metallurgical meaning, it will act on it in about
        one case in six, and neither accuracy, per-class metrics, the heatmap, nor the
        obvious hold-out test can see it.* Latent until a line re-labels, a microscope is
        swapped, or sample prep changes — which is why this belongs in periodic
        revalidation, not a one-time sign-off.
        """


SPURIOUS_MARKDOWN = _spurious_markdown()

example_dir = "examples"
example_paths = []
if os.path.isdir(example_dir):
    example_paths = [
        [os.path.join(example_dir, f)]
        for f in sorted(os.listdir(example_dir))
        if f.lower().endswith((".jpg", ".jpeg", ".png"))
    ]

# Pin the page to light mode. The embedded result figures are matplotlib defaults with
# white backgrounds, which read as glaring rectangles inside a dark shell; Gradio has no
# "always light" flag, so we set its own __theme param on load and let the page settle.
FORCE_LIGHT_THEME = """
function() {
    const url = new URL(window.location);
    if (url.searchParams.get('__theme') !== 'light') {
        url.searchParams.set('__theme', 'light');
        window.location.replace(url.href);
    }
}
"""

# The tab bar is how a visitor discovers that tabs 2 and 3 exist at all, and those
# carry the whole contribution. Gradio's default marks the active tab by colour alone
# -- same weight, no fill, no border -- which is a weak cue and useless to anyone with
# red-green colour deficiency; and its container is flex-wrap:nowrap + overflow:hidden,
# so on a 375px phone tab 3 lands offscreen with no way to scroll to it. Below: three
# independent active-state signals (weight, fill, underline) and a scrollable bar.
# Targets .tab-container and aria-selected rather than Gradio's svelte-* hash, which is
# regenerated on every SDK rebuild and would silently stop matching.
TAB_CSS = """
.tab-container { gap: 6px; overflow-x: auto; scrollbar-width: none; }
.tab-container::-webkit-scrollbar { display: none; }
.tab-container button[role="tab"] {
    flex: 0 0 auto;
    white-space: nowrap;
    font-size: 15px;
    font-weight: 500;
    padding: 11px 20px;
    color: #52525b;
    background: #f4f4f5;
    border: 1px solid #e5e7eb;
    border-radius: 8px 8px 0 0;
}
.tab-container button[role="tab"]:hover { background: #ececee; color: #27272a; }
.tab-container button[role="tab"][aria-selected="true"] {
    font-weight: 600;
    color: #9a3412;
    background: #ffffff;
    border-bottom: 3px solid #f97316;
}
@media (max-width: 640px) {
    .tab-container button[role="tab"] { font-size: 13px; padding: 9px 12px; }
}
"""

with gr.Blocks(title="Explainable & Auditable Alloy Defect Detection",
               css=TAB_CSS, js=FORCE_LIGHT_THEME) as demo:
    gr.Markdown(
        """
        # Explainable & Auditable Alloy Defect Detection
        Group 8's submission for Problem #10, MIT Professional Education "Applied
        AI for Materials Discovery." A steel-surface defect classifier (NEU-DET, 6
        defect classes) plus the interpretability and statistical-guarantee tools a
        plant manager or auditor would actually need before acting on it.

        **Read the tabs in order** — they are an argument, not a menu.
        The model works and explains itself; that explanation can be hijacked by a
        meaningless artifact; and while it does pass the standard faithfulness test,
        that test turns out to be **sensitive to how you split the data**. What we would
        put in front of an auditor is none of those — it is a statistical
        guarantee, plus the lesson that audit tools need auditing before you trust them.
        """
    )
    with gr.Tab("1 - Classify & Explain"):
        gr.Markdown(
            "### A working model that explains itself\n\n"
            "Upload a micrograph (or pick an example below). The model predicts the "
            "most likely defect class; the heatmap shows which pixels drove that "
            "prediction (**Grad-CAM**, implemented from first principles with PyTorch "
            "hooks, following the course's own notebook).\n\n"
            "*The classifier itself is deliberately ordinary* — a compact 164k-parameter "
            "CNN, trained by the same recipe as the notebook: a stratified 70/10/10/10 split "
            "where the test images are never used to fit weights or pick a checkpoint. "
            "The notebook's run scores 94.1% on that test split; the separately trained "
            "checkpoint served here scores 95.6% on the same split — same recipe and "
            "same split, different training run. That is the "
            "point: the question is not whether a CNN can score well on NEU-DET (it can), but "
            "whether the explanation sitting on top of it means anything. Tabs 2 and 3 test "
            "exactly that."
        )
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(type="pil", label="Micrograph")
                run_button = gr.Button("Classify", variant="primary")
                gr.Examples(examples=example_paths, inputs=image_input)
            with gr.Column():
                overlay_output = gr.Image(label="Grad-CAM overlay (predicted class)")
                label_output = gr.Label(label="Predicted class probabilities", num_top_classes=6)
        run_button.click(predict, inputs=image_input, outputs=[overlay_output, label_output])
        image_input.change(predict, inputs=image_input, outputs=[overlay_output, label_output])

    with gr.Tab("2 - Can We Fool the Model?"):
        gr.Markdown(SPURIOUS_SETUP)
        gr.Image(
            value="assets/spurious-correlation.png",
            label="Grad-CAM locks onto a planted batch stamp, whether or not the prediction flips",
            interactive=False,
        )
        gr.Markdown(SPURIOUS_MARKDOWN)

    with gr.Tab("3 - What We Show an Auditor"):
        gr.Markdown(
            """
            ### First: can we even tell whether the heatmap is faithful?

            Tab 2 showed the explanation can be captured by an artifact we planted
            ourselves. Fair objection: real data has no planted stamp. So we tested the
            heatmap directly on clean images, using **the course's own deletion test**
            — delete the pixels Grad-CAM ranks most important and measure how much
            the prediction suffers, versus deleting the same number of random pixels.
            """
        )
        gr.Image(
            value="assets/deletion-curve.png",
            label="Deletion test: Grad-CAM-ranked pixels versus random pixels (precomputed in the notebook)",
            interactive=False,
        )
        gr.Markdown(
            """
            **It passes cleanly.** At 30% of pixels replaced, Grad-CAM-guided deletion has
            cost **0.898** of predicted probability; random deletion has cost **0.016**. The
            pixels the heatmap points to are the pixels the prediction actually depends on.

            > **A note on method.** An earlier version of this project used an 85/15 split
            > and got the *opposite* verdict on this same test (0.546 versus 0.718). We
            > therefore ran it properly — two split designs, three seeds each, six
            > independent runs — and the separation was perfect: 3/3 unfavourable under
            > 85/15, 3/3 favourable under the 70/10/10/10 split used here. Notably it is the
            > *random* baseline that moves, not the attribution. The full analysis is in the
            > notebook. We report the clean-split result because that design never lets test
            > images influence weights or checkpoint selection — but the sensitivity is
            > why we treat this curve as supporting evidence rather than proof, and why the
            > guarantee below is what we would actually put in front of an auditor.

            ---

            ### So what survives? Conformal Risk Control

            **Not covered anywhere in the course** — taken from Shen & Liu,
            [arXiv:2504.17721](https://arxiv.org/abs/2504.17721) (2025). It wraps around
            *any* existing detector and needs no retraining.

            **Why not just calibrate the confidence scores?** Calibration says *"when I
            output 0.9, I'm right about 90% of the time"* — a description of average
            past behaviour, true only while conditions hold, and guaranteed by nothing.
            Conformal risk control says something categorically different: *"the expected
            error rate on future batches is **at most** α"* — a bound that is
            **derived, not fitted**. To an auditor those are not the same kind of evidence.
            (It carries preconditions, and we examine ours in the limits below rather than
            waving them through.)

            **How it works, in one paragraph.** Hold back a set of images the model never
            trained on. Sweep a confidence bar across them and find the least strict bar
            that still keeps the error rate under your chosen budget α. Provided
            future images come from the same process as those held-back ones
            (*exchangeability*), a future image is just another draw from the same pool
            — so the bar that worked on them must work in expectation on it too.
            The guarantee is arithmetic, not optimism.

            **It was validated, not just asserted:** empirical error stayed at or below
            target at **six of the nine** risk levels tested (α = 0.1 through 0.9).
            The three that came in over — α = 0.10, 0.20 and 0.90 — miss by
            **1.12, 0.55 and 0.83 standard errors** on 180 test images. The bound is on
            *expected* error, so realisations scatter either side of the line; near-misses in
            both directions are what a correctly behaving bound looks like. An earlier version
            of this project reported nine out of nine, which looked stronger and was weaker
            — that model had been checkpoint-selected on the very images it was scored
            against.
            """
        )
        gr.Image(
            value="assets/conformal-guarantee.png",
            label="Conformal Risk Control: empirical FDR against target at nine risk levels; under target at six, over by ~1 standard error at three (precomputed in the notebook)",
            interactive=False,
        )
        gr.Markdown(
            """
            #### What the plant actually gets

            For each micrograph the method returns a **prediction set** — the defect
            classes clearing the calibrated bar. Its *size* is the instruction:

            | Prediction set | Meaning | Action |
            |---|---|---|
            | Exactly one class | Confident and bounded | **Act** — route the batch |
            | Two or more | Genuinely ambiguous at this risk level | **Escalate**, candidates attached |
            | Empty | Nothing clears the bar | **Abstain** — the model declines |

            The plant picks α. That single dial trades throughput against escalation
            volume. Pick it at the **strict end**: α = 0.10 abstains on 12.8% of images,
            α = 0.90 on 91.7% — within budget and operationally useless. **This, not
            the accuracy number and not the heatmap, is what we would put in front of an
            auditor.**

            #### Four honest limits

            - **The bound is an average, not a per-image promise.** It constrains the error
              rate across future batches. It does *not* say this particular micrograph has
              a ≤α chance of being wrong, and must not be sold that way.
            - **Exchangeability is the whole vulnerability.** A new production line, a
              replaced microscope, or changed sample prep breaks it *silently* — no
              error is raised, the number just quietly stops meaning what you think.
              Recalibrate on data from the new source.
            - **At lenient α the metric is dominated by abstentions.** We adapted a
              method built for pixel-level segmentation, where a prediction set spans
              thousands of pixels and is essentially never empty. With six classes, empty
              sets are common at strict bars — 91.7% of images at the most lenient
              α — and an empty set is scored as a full error. So at that end the
              reported "false discovery rate" is very nearly a synonym for "abstention rate."
              The useful operating range here is the **low-α end**, where empty sets are
              far rarer (12.8%).
            - **A precondition of the theorem is not strictly met.** Conformal Risk Control
              requires the loss to be *monotone non-increasing* — as you flag more, the
              risk must never climb back. False-negative rate satisfies that; **false
              discovery rate does not.** Ours is U-shaped: it falls from 0.949 to a minimum of
              0.029 around λ ≈ 0.461, then climbs back to 0.755 as the
              flagged set fills with false positives — rising at 80 of 399 grid steps.
              Every threshold we actually selected (0.004–0.219) sits on the *descending*
              branch, where the condition effectively holds; the first rise is not until
              λ = 0.486. Capping the search grid at λ ≤ 0.48
              would satisfy the precondition outright **without changing a single reported
              number**. But that is a property of this data, not a proof. **It is why the chart
              above is load-bearing rather than decorative:** with a precondition in question,
              checking the bound empirically on held-out data is doing real work.
            """
        )

    gr.Markdown(
        "---\n**The full evidence behind every number on this page is in this Space's "
        "[`notebook/`](https://huggingface.co/spaces/joenathan/"
        "mit-group8-explainable-defect-detection/tree/main/notebook) folder** — the "
        "[Jupyter notebook](https://huggingface.co/spaces/joenathan/"
        "mit-group8-explainable-defect-detection/blob/main/notebook/"
        "mit-group8-explainable-defect-detection.ipynb) with all outputs, and a "
        "[plain-language walkthrough](https://huggingface.co/spaces/joenathan/"
        "mit-group8-explainable-defect-detection/blob/main/notebook/"
        "mit-group8-notebook-walkthrough.md) of what each cell does and why. The charts "
        "above are precomputed there; the notebook is where you can check them.\n\n"
        "Group 8 · MIT Professional Education, Applied AI for Materials "
        "Discovery, July 2026."
    )

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