"""WhatBird — small + large model bird identification.

Pipeline:
    photo --> [classify] yolo26x ONNX specialist --> top-5 of 1,532 species
          --> [reason]   MiniCPM-V 4.6 (1B)       --> re-rank + explain shortlist
          --> verdict + confidence bars + field-mark reasoning + saliency map

The classifier runs anywhere (ONNX Runtime, CPU). The vision-language model
switches on automatically wherever a GPU can serve it (HF ZeroGPU Space, local
CUDA box); elsewhere a templated fallback keeps the app functional. No
configuration needed (WHATBIRD_DESCRIBER / WHATBIRD_MODEL_ID override it if set).
"""

from __future__ import annotations

import gradio as gr
from PIL import Image

from whatbird.classifier import BirdClassifier
from whatbird.describer import get_describer
from whatbird.saliency import occlusion_heatmap

TOPK = 5

classifier = BirdClassifier()
describer = get_describer()


def identify(image: Image.Image, show_heatmap: bool):
    hide = gr.update(value=None, visible=False)
    if image is None:
        return {}, "Upload or capture a photo of a bird to begin.", "", hide, hide
    candidates = classifier.predict(image, topk=TOPK)
    # A few species exist under two raw labels (dataset-merge leftovers, e.g.
    # Harris/Harriss Sparrow) — keep the higher-confidence entry per name.
    label_scores: dict[str, float] = {}
    for c in candidates:
        label_scores[c.name] = max(label_scores.get(c.name, 0.0), c.confidence)

    verdict = describer.describe(image, candidates)
    note = ""
    if not verdict.in_shortlist:
        note = (
            "\n\n> **Not in the classifier's shortlist.** The classifier was "
            "unsure, so the vision-language model identified this species directly "
            "from the photo."
        )
    explanation = f"### {verdict.species}\n\n{verdict.explanation}{note}"
    links = "  ·  ".join(
        dict.fromkeys(f"[{c.name}]({c.wiki})" for c in candidates)
    )
    mode = (
        f"explained by {verdict.source}"
        if verdict.source != "stub"
        else "fast mode — no vision-language model on this hardware"
    )
    refs = f"**Look it up:** {links}\n\n<sub>{mode}</sub>"

    # Explain the *verdict's* species, not blindly the classifier's top-1; on an
    # off-list pick (verdict.label is None) the classifier never saw the species,
    # so a heatmap of its wrong guess would mislead — hide it.
    heatmap, caption = hide, gr.update(visible=False)
    if show_heatmap and verdict.label is not None:
        target = next(c for c in candidates if c.label == verdict.label)
        heatmap = gr.update(
            value=occlusion_heatmap(classifier, image, target.index), visible=True
        )
        caption = gr.update(visible=True)
    return label_scores, explanation, refs, heatmap, caption


THEME = gr.themes.Soft(primary_hue="emerald", secondary_hue="teal")

with gr.Blocks(title="WhatBird") as demo:
    gr.Markdown(
        "# 🐦 WhatBird\n"
        "**Two-stage bird identification.** A fast on-device classifier "
        "(yolo26x, 1,532 species) shortlists the candidates; a compact vision-language "
        "model (MiniCPM-V 4.6, 1B) then looks at your photo to confirm the "
        "species and explain *why* from its field marks."
    )
    with gr.Row():
        with gr.Column(scale=1):
            inp = gr.Image(type="pil", sources=["upload", "webcam"], label="Bird photo")
            show_heatmap = gr.Checkbox(value=True, label="Show where the classifier looked")
            btn = gr.Button("Identify", variant="primary")
            gr.Examples(
                examples=[
                    ["samples/sample_1.jpg"],
                    ["samples/sample_2.jpg"],
                    ["samples/sample_3.jpg"],
                    ["samples/robin.jpg"],  # European Robin — now in-domain (added via iNaturalist)
                ],
                inputs=inp,
            )
        with gr.Column(scale=1):
            out_label = gr.Label(num_top_classes=TOPK, label="Top candidates")
            out_md = gr.Markdown()
            out_refs = gr.Markdown()
            out_heatmap = gr.Image(label="Saliency", type="pil", visible=False)
            out_caption = gr.Markdown(
                "<sub>Brighter = the regions the classifier relied on most "
                "(occlusion saliency).</sub>",
                visible=False,
            )

    outputs = [out_label, out_md, out_refs, out_heatmap, out_caption]
    btn.click(identify, inputs=[inp, show_heatmap], outputs=outputs)


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