"""Constructors for MoPET models and pretrained-weight loading.

``create_model`` builds a MoPET model on a timm backbone with the paper's default
expert configuration; ``load_pretrained_weights`` fetches published checkpoints
from the HuggingFace Hub. The MedMNIST class-count table is kept here so the
package can size its heads without importing ``medmnist``.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import cast

import timm
import torch
from torch import nn

from .model import MoPET

logger = logging.getLogger(__name__)

#: Friendly backbone name -> timm model id (all ViT-Base/16).
BACKBONES: dict[str, str] = {
    "dinov3": "vit_base_patch16_dinov3.lvd1689m",
    "dino": "vit_base_patch16_224.dino",
    "clip": "vit_base_patch16_clip_224",
}

#: Number of classes per MedMNIST+ 2D dataset (used to size the per-dataset heads).
MEDMNIST_NUM_CLASSES: dict[str, int] = {
    "BloodMNIST": 8,
    "BreastMNIST": 2,
    "ChestMNIST": 14,
    "DermaMNIST": 7,
    "OCTMNIST": 4,
    "OrganAMNIST": 11,
    "OrganCMNIST": 11,
    "OrganSMNIST": 11,
    "PathMNIST": 9,
    "PneumoniaMNIST": 2,
    "RetinaMNIST": 5,
    "TissueMNIST": 8,
}

#: Default heterogeneous expert pool (the MoPET configuration from the paper).
DEFAULT_EXPERT_COUNTS: dict[str, int] = {"LoRA": 20, "BOFT": 12}
DEFAULT_TOP_K: int = 12
DEFAULT_EXPERT_KWARGS: dict[str, dict[str, object]] = {
    "LoRA": {"r": 8, "lora_alpha": 8},
    "BOFT": {"boft_block_size": 8, "boft_n_butterfly_factor": 1},
    "FourierFT": {"n_frequency": 1000},
}


@dataclass(frozen=True)
class PublishedModel:
    """A released MoPET checkpoint: its HuggingFace repo and the exact head layout.

    ``datasets`` is ordered: index i is the dataset id of the i-th classification head,
    so it must match the order the checkpoint was trained with.
    """

    repo_id: str
    backbone: str
    datasets: tuple[str, ...]


#: Released MoPET models (the paper's headline checkpoints). ``create_model(weights=<name>)``
#: and ``load_pretrained_weights(model, <name>)`` resolve these and pull from the Hub.
PUBLISHED_MODELS: dict[str, PublishedModel] = {
    # 4-dataset unified model (Table 2) — also the Breast-booster (Table 3, same pool).
    "unified": PublishedModel(
        "sdoerrich97/mopet_dinov3_unified_blood_breast_derma_path",
        "dinov3",
        ("BloodMNIST", "BreastMNIST", "DermaMNIST", "PathMNIST"),
    ),
    # Auxiliary-booster models (Table 3): a target co-trained with a hand-picked pool.
    "booster-retina": PublishedModel(
        "sdoerrich97/mopet_dinov3_booster_retina_breast_blood_retina_path_organa",
        "dinov3",
        ("BreastMNIST", "BloodMNIST", "RetinaMNIST", "PathMNIST", "OrganAMNIST"),
    ),
    "booster-derma": PublishedModel(
        "sdoerrich97/mopet_dinov3_booster_derma_derma_blood_oct_organs",
        "dinov3",
        ("DermaMNIST", "BloodMNIST", "OCTMNIST", "OrganSMNIST"),
    ),
}


def list_pretrained() -> dict[str, dict[str, object]]:
    """List the released MoPET checkpoints and how to load them.

    Returns:
        A mapping from each variant name (the string passed as ``create_model(weights=...)``)
        to its ``backbone``, ordered ``datasets`` (index i is the dataset id of head i), and
        HuggingFace ``repo_id``. Use it to discover the available weights, e.g.::

            import mopet
            for name, info in mopet.list_pretrained().items():
                print(name, info["datasets"])
    """
    return {
        name: {
            "backbone": pub.backbone,
            "datasets": list(pub.datasets),
            "repo_id": pub.repo_id,
        }
        for name, pub in PUBLISHED_MODELS.items()
    }


def _resolve_backbone(backbone: str) -> str:
    """Map a friendly backbone name to its timm id (pass-through if already an id)."""
    return BACKBONES.get(backbone, backbone)


def _resolve_num_classes(datasets: list[str] | None, num_classes: list[int] | None) -> list[int]:
    """Resolve the per-dataset class counts from dataset names or an explicit list."""
    if num_classes is not None:
        return num_classes
    if datasets is not None:
        try:
            return [MEDMNIST_NUM_CLASSES[d] for d in datasets]
        except KeyError as exc:  # pragma: no cover - defensive
            raise KeyError(f"Unknown MedMNIST dataset: {exc.args[0]!r}") from exc
    raise ValueError("Provide either `datasets` or `num_classes` to size the heads.")


def create_model(
    backbone: str = "dinov3",
    datasets: list[str] | None = None,
    num_classes: list[int] | None = None,
    pretrained_backbone: bool = True,
    weights: str | None = None,
    expert_counts: dict[str, int] | None = None,
    top_k: int = DEFAULT_TOP_K,
    expert_kwargs: dict[str, dict[str, object]] | None = None,
    controller_noise: bool = False,
    map_location: str = "cpu",
) -> MoPET:
    """Build a MoPET model.

    Args:
        backbone: Friendly name (``"dinov3"``/``"dino"``/``"clip"``) or a timm id.
        datasets: MedMNIST dataset names defining the multi-task heads (order matters).
        num_classes: Explicit class counts per head; overrides ``datasets``.
        pretrained_backbone: Load timm pretrained backbone weights.
        weights: Either a published-variant name (a key of ``PUBLISHED_MODELS``, e.g.
            ``"unified"``) to download from the Hub, or a path to a local MoPET
            checkpoint. A variant name also fixes the backbone and head layout.
        expert_counts: Experts per family; defaults to the paper's ``{LoRA:20, BOFT:12}``.
        top_k: Experts activated per token.
        expert_kwargs: Per-family peft keyword arguments; defaults to the paper's.
        controller_noise: Whether routers add exploration noise in training.
        map_location: Device mapping used when loading ``weights``.

    Returns:
        The constructed :class:`~mopet.model.MoPET`.
    """
    published = PUBLISHED_MODELS.get(weights) if weights is not None else None
    if published is not None:
        # A released variant fixes the backbone and head layout. Published checkpoints carry
        # only the trainable parameters (adapters, router, heads); the frozen backbone is
        # reconstructed from the timm pretrained weights, so keep `pretrained_backbone=True`.
        backbone = published.backbone
        datasets = list(published.datasets)
        num_classes = None
        pretrained_backbone = True

    timm_id = _resolve_backbone(backbone)
    heads = _resolve_num_classes(datasets, num_classes)
    backbone_module = cast(
        nn.Module, timm.create_model(timm_id, pretrained=pretrained_backbone, num_classes=0)
    )

    model = MoPET(
        backbone=backbone_module,
        num_classes=heads,
        expert_counts=expert_counts or dict(DEFAULT_EXPERT_COUNTS),
        top_k=top_k,
        expert_kwargs=expert_kwargs or {k: dict(v) for k, v in DEFAULT_EXPERT_KWARGS.items()},
        controller_noise=controller_noise,
    )
    if published is not None:
        _load_into(model, _load_state_dict(_download_variant(published), map_location=map_location))
    elif weights is not None:
        _load_into(model, _load_state_dict(weights, map_location=map_location))
    return model


def load_pretrained_weights(model: MoPET, variant: str, map_location: str = "cpu") -> MoPET:
    """Download and load a published MoPET checkpoint from the HuggingFace Hub.

    Args:
        model: A MoPET model whose head layout matches ``variant`` (build it with the same
            ``datasets``/``backbone``, e.g. via ``create_model(weights=variant)``).
        variant: A key of ``PUBLISHED_MODELS`` (e.g. ``"unified"``, ``"booster-retina"``).
        map_location: Device mapping for the loaded tensors.

    Returns:
        ``model`` with the checkpoint loaded in place.
    """
    if variant not in PUBLISHED_MODELS:
        raise KeyError(
            f"No published weights for {variant!r}. Available: {sorted(PUBLISHED_MODELS)}."
        )
    _load_into(model, _load_state_dict(_download_variant(PUBLISHED_MODELS[variant]), map_location))
    return model


def _download_variant(published: PublishedModel) -> str:
    """Download a published variant's weights from the Hub, preferring safetensors."""
    from huggingface_hub import hf_hub_download

    try:
        return hf_hub_download(published.repo_id, filename="model.safetensors")
    except Exception:  # noqa: BLE001 - fall back to a torch checkpoint
        return hf_hub_download(published.repo_id, filename="model.pth")


def _load_state_dict(path: str, map_location: str = "cpu") -> dict[str, torch.Tensor]:
    """Load a state dict from a ``.safetensors`` or torch checkpoint file."""
    if path.endswith(".safetensors"):
        from safetensors.torch import load_file

        return load_file(path, device=map_location)
    obj = torch.load(path, map_location=map_location, weights_only=False)
    state_dict = obj.get("state_dict", obj) if isinstance(obj, dict) else obj
    return {k.removeprefix("module."): v for k, v in state_dict.items()}


def _load_into(model: MoPET, state_dict: dict[str, torch.Tensor]) -> None:
    """Load trainable parameters into ``model``, tolerating the frozen backbone gap."""
    missing, unexpected = model.load_state_dict(state_dict, strict=False)
    if unexpected:
        logger.warning("Unexpected keys when loading MoPET weights: %s", unexpected[:8])
    logger.info(
        "Loaded MoPET weights (%d missing, %d unexpected keys).", len(missing), len(unexpected)
    )