"""MoPET: Parameter-Efficient Mixture-of-Experts for Unified Medical Image Classification.

A Gradio demo that loads the pretrained MoPET ``unified`` checkpoint and lets
visitors upload a medical image, pick one of four MedMNIST+ tasks, and receive
a classification from the model's per-dataset head.
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch / any CUDA-touching import

import gradio as gr
import timm
import torch
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms

from mopet import create_model
from mopet._factory import BACKBONES, PUBLISHED_MODELS

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

VARIANT = "unified"
PUBLISHED = PUBLISHED_MODELS[VARIANT]
DATASET_NAMES = list(PUBLISHED.datasets)  # head order = dataset id

# MedMNIST+ class labels (from the official INFO table)
CLASS_LABELS: dict[str, list[str]] = {
    "BloodMNIST": [
        "basophil",
        "eosinophil",
        "erythroblast",
        "immature granulocytes",
        "lymphocyte",
        "monocyte",
        "neutrophil",
        "platelet",
    ],
    "BreastMNIST": ["malignant", "normal, benign"],
    "DermaMNIST": [
        "actinic keratoses",
        "basal cell carcinoma",
        "benign keratosis-like lesions",
        "dermatofibroma",
        "melanoma",
        "melanocytic nevi",
        "vascular lesions",
    ],
    "PathMNIST": [
        "adipose",
        "background",
        "debris",
        "lymphocytes",
        "mucus",
        "smooth muscle",
        "normal colon mucosa",
        "cancer-associated stroma",
        "colorectal adenocarcinoma epithelium",
    ],
}

# DINOv3 backbone normalization stats (from timm pretrained config)
_tim_id = BACKBONES[PUBLISHED.backbone]
_cfg = timm.get_pretrained_cfg(_tim_id)
_MEAN = tuple(_cfg.mean)
_STD = tuple(_cfg.std)

# Preprocessing: ToTensor → Normalize → Pad to 224 → Resize to 256
# (matches the paper's `build_transform` for DINOv3 at resolution 224)
_padding = max(0, 224 - 224)  # 0 when loading at 224; images are resized anyway
_pad_l = _pad_t = _padding // 2
_pad_r = _padding - _pad_l
_pad_b = _padding - _pad_t

TRANSFORM = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=_MEAN, std=_STD),
    transforms.Pad((_pad_l, _pad_t, _pad_r, _pad_b), fill=0, padding_mode="constant"),
    transforms.Resize((256, 256)),
])

# ---------------------------------------------------------------------------
# Model loading (module scope, eager .to("cuda"))
# ---------------------------------------------------------------------------

print(f"Loading MoPET '{VARIANT}' (backbone={PUBLISHED.backbone}, datasets={DATASET_NAMES}) ...")
model = create_model(weights=VARIANT, map_location="cpu").eval()
model = model.to("cuda")
print("Model loaded and moved to CUDA.")


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------

@spaces.GPU(duration=60)
def classify(image: Image.Image, dataset_name: str) -> dict:
    """Classify a medical image using the MoPET unified model.

    Args:
        image: An RGB medical image (blood cell, breast ultrasound, dermoscopy,
            or colon pathology histology).
        dataset_name: Which MedMNIST+ task head to use for classification.

    Returns:
        A label-probability dictionary for the selected task.
    """
    if image is None:
        return {label: 0.0 for label in CLASS_LABELS[dataset_name]}

    image = image.convert("RGB")
    tensor = TRANSFORM(image).unsqueeze(0).to("cuda")

    dataset_id = DATASET_NAMES.index(dataset_name)
    dataset_ids = torch.tensor([dataset_id], device="cuda")

    with torch.no_grad():
        logits = model(tensor, dataset_ids)
        probs = F.softmax(logits, dim=-1).squeeze(0).cpu()

    labels = CLASS_LABELS[dataset_name]
    num_classes = len(labels)
    # The model pads to C_max; only the first ``num_classes`` entries are real
    return {labels[i]: float(probs[i]) for i in range(num_classes)}


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

CSS = """
#col-container { max-width: 900px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    gr.Markdown(
        "# MoPET: Parameter-Efficient Mixture-of-Experts for Unified Medical Image Classification\n"
        "Upload a medical image and select a task. The model routes the image through a "
        "frozen DINOv3 backbone with a sparse mixture-of-experts adapter pool and a "
        "per-task classification head.\n\n"
        "[Paper](https://arxiv.org/abs/2607.29462) | "
        "[GitHub](https://github.com/sdoerrich97/mopet) | "
        "[Weights](https://huggingface.co/sdoerrich97/mopet_dinov3_unified_blood_breast_derma_path)"
    )

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            image_input = gr.Image(type="pil", label="Medical Image", scale=3)
            dataset_selector = gr.Dropdown(
                choices=DATASET_NAMES,
                value=DATASET_NAMES[0],
                label="Task / Dataset Head",
                scale=1,
            )

        run_btn = gr.Button("Classify", variant="primary")

        label_output = gr.Label(num_top_classes=5, label="Classification")

    run_btn.click(
        fn=classify,
        inputs=[image_input, dataset_selector],
        outputs=label_output,
        api_name="classify",
    )

    gr.Examples(
        examples=[
            ["bloodmnist_sample.png", "BloodMNIST"],
            ["breastmnist_sample.png", "BreastMNIST"],
            ["dermamnist_sample.png", "DermaMNIST"],
            ["pathmnist_sample.png", "PathMNIST"],
        ],
        inputs=[image_input, dataset_selector],
        outputs=label_output,
        fn=classify,
        cache_examples=True,
        cache_mode="lazy",
    )

demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)