"""The MoPET model: a frozen foundation backbone adapted by a routed PEFT expert pool.

``MoPET`` freezes a timm vision-transformer backbone, replaces its attention
``qkv`` projections with sparse mixture-of-experts adapters (see :mod:`mopet._moe`),
and attaches one classification head per dataset so a single network serves many
heterogeneous medical-image classification tasks at once.
"""

from __future__ import annotations

import logging
from typing import cast

import torch
from torch import nn

from ._moe import MoEModule, apply_moe_peft

logger = logging.getLogger(__name__)


class MultiTaskClassifier(nn.Module):
    """One linear classification head per dataset with padded, per-sample routing.

    Args:
        num_classes: Number of classes for each dataset, in a fixed order; the
            index into this list is the dataset identifier used at forward time.
        input_dim: Dimension of the shared backbone feature.
    """

    def __init__(self, num_classes: list[int], input_dim: int) -> None:
        super().__init__()
        self.classifiers = nn.ModuleList(
            nn.Linear(in_features=input_dim, out_features=n) for n in num_classes
        )

    def forward(self, x: torch.Tensor, dataset_ids: torch.Tensor) -> torch.Tensor:
        """Route each sample to its dataset head.

        Args:
            x: Shared features of shape ``(B, D)``.
            dataset_ids: Dataset index per sample, shape ``(B,)``.

        Returns:
            Logits of shape ``(B, C_max)`` where ``C_max`` is the largest class
            count present in the batch; unused entries are padded with ``-1e9``.
        """
        b = x.size(0)
        present: list[int] = torch.unique(dataset_ids).tolist()
        heads = [cast(nn.Linear, self.classifiers[task]) for task in present]
        max_classes = max(head.out_features for head in heads)

        output = torch.full((b, max_classes), fill_value=-1e9, device=x.device, dtype=x.dtype)
        for task, head in zip(present, heads, strict=True):
            idx = (dataset_ids == task).nonzero(as_tuple=True)[0]
            logits = head(x[idx])
            output[idx, : logits.size(1)] = logits
        return output


class MoPET(nn.Module):
    """Frozen backbone + routed PEFT experts + per-dataset heads.

    Args:
        backbone: A timm vision transformer providing ``(B, D)`` pooled features
            once its own head is removed. Frozen in place.
        num_classes: Class count per dataset (defines the multi-task heads).
        expert_counts: Number of experts per family, e.g. ``{"LoRA": 20, "BOFT": 12}``.
        top_k: Number of experts activated per token.
        expert_kwargs: Per-family peft keyword arguments.
        controller_noise: Whether routers add exploration noise in training.
    """

    def __init__(
        self,
        backbone: nn.Module,
        num_classes: list[int],
        expert_counts: dict[str, int],
        top_k: int,
        expert_kwargs: dict[str, dict[str, object]],
        controller_noise: bool = False,
    ) -> None:
        super().__init__()
        self.pretrained_cfg = getattr(backbone, "pretrained_cfg", None)

        backbone.head = nn.Identity()
        for param in backbone.parameters():
            param.requires_grad = False

        self.backbone = apply_moe_peft(
            backbone,
            expert_counts=expert_counts,
            top_k=top_k,
            expert_kwargs=expert_kwargs,
            controller_noise=controller_noise,
        )
        input_dim = int(cast(int, backbone.num_features))
        self.head = MultiTaskClassifier(num_classes=num_classes, input_dim=input_dim)

    def forward(self, x: torch.Tensor, dataset_ids: torch.Tensor) -> torch.Tensor:
        """Classify a batch of images tagged with their dataset ids.

        Args:
            x: Input images of shape ``(B, C, H, W)``.
            dataset_ids: Dataset index per sample, shape ``(B,)``.

        Returns:
            Padded per-dataset logits of shape ``(B, C_max)``.
        """
        features = self.backbone(x)
        return self.head(features, dataset_ids)

    def get_aux_loss(self) -> torch.Tensor | float:
        """Sum the load-balancing loss over all routed layers from the last forward."""
        aux_loss: torch.Tensor | float = 0.0
        for module in self.backbone.modules():
            if isinstance(module, MoEModule):
                aux_loss = aux_loss + module.get_aux_loss()
        return aux_loss