"""Sparse mixture-of-experts over parameter-efficient adapters.

The core of *MoPET*: a learnable top-k router dispatches each token to a small
subset of a heterogeneous pool of PEFT experts (LoRA / BOFT / FourierFT) that
wrap a single frozen projection. Each expert returns the full projection output
(frozen base plus its low-rank delta); because the router weights are a softmax
over the selected experts, their sum reduces to the frozen projection plus a
convex combination of the active experts' deltas.

This module is deliberately free of any dataset or configuration framework
dependency: expert counts, the routing width, and the per-family adapter
hyperparameters are passed in explicitly.
"""

from __future__ import annotations

import logging

import torch
import torch.nn.functional as F
from peft.tuners.boft.layer import Linear as BOFTLinear
from peft.tuners.fourierft import FourierFTLinear
from peft.tuners.lora import Linear as LoRALinear
from torch import nn

logger = logging.getLogger(__name__)

#: Maps an expert-family name to the peft layer class that implements it. The
#: internal ``peft`` layer classes are used directly (as in the original thesis
#: code); this is why ``peft`` is pinned exactly.
_EXPERT_CLASSES: dict[str, type[nn.Module]] = {
    "LoRA": LoRALinear,
    "BOFT": BOFTLinear,
    "FourierFT": FourierFTLinear,
}


class LinearTopKGating(nn.Module):
    """Linear router producing per-expert routing logits for each token.

    Args:
        input_dim: Token feature dimension.
        num_experts: Size of the expert pool to route over.
        noisy: If ``True``, additive standard-normal noise is applied to the
            logits during training only, to encourage exploration.
    """

    def __init__(self, input_dim: int, num_experts: int, noisy: bool = False) -> None:
        super().__init__()
        self.num_experts = num_experts
        self.noisy = noisy
        # g(x) = W x, no bias.
        self.gate = nn.Linear(input_dim, num_experts, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Return unnormalized routing logits of shape ``(..., num_experts)``."""
        logits = self.gate(x)
        if self.noisy and self.training:
            logits = logits + torch.randn_like(logits)
        return logits


class MoEModule(nn.Module):
    """Replaces a single frozen projection with a routed pool of PEFT experts.

    Injected in place of an attention ``qkv`` projection. The wrapped frozen
    linear is shared as the base layer of every expert, so each expert output is
    ``base(x) + delta_i(x)`` and the routed, softmax-weighted sum is
    ``base(x) + sum_i w_i * delta_i(x)`` over the active experts.

    Args:
        base_linear: The frozen projection to adapt (e.g. attention ``qkv``).
        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 keyword arguments forwarded to the peft layer,
            e.g. ``{"LoRA": {"r": 8, "lora_alpha": 8}, "BOFT": {...}}``.
        controller_noise: Whether the router adds exploration noise in training.
        adapter_name: Adapter name handed to the peft layers.
    """

    def __init__(
        self,
        base_linear: nn.Linear,
        expert_counts: dict[str, int],
        top_k: int,
        expert_kwargs: dict[str, dict[str, object]],
        controller_noise: bool = False,
        adapter_name: str = "default",
    ) -> None:
        super().__init__()
        self.top_k = top_k
        in_dim = base_linear.in_features
        self.out_features = base_linear.out_features

        self.experts = nn.ModuleList()
        for family, count in expert_counts.items():
            if family not in _EXPERT_CLASSES:
                raise ValueError(f"Unsupported expert family: {family!r}")
            expert_cls = _EXPERT_CLASSES[family]
            kwargs = expert_kwargs.get(family, {})
            for _ in range(count):
                self.experts.append(
                    expert_cls(base_layer=base_linear, adapter_name=adapter_name, **kwargs)
                )

        self.num_experts = len(self.experts)
        self.controller = LinearTopKGating(
            input_dim=in_dim, num_experts=self.num_experts, noisy=controller_noise
        )
        self._aux_loss: torch.Tensor | float = 0.0

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Route tokens through the top-k experts.

        Args:
            x: Token features of shape ``(B, T, D)``.

        Returns:
            Adapted projection output of shape ``(B, T, out_features)``.
        """
        b, t, _ = x.shape
        logits = self.controller(x)  # (B, T, E)

        topk_vals, topk_idx = torch.topk(logits, self.top_k, dim=-1)  # (B, T, k)
        topk_weights = F.softmax(topk_vals, dim=-1)  # (B, T, k)

        output = torch.zeros(b, t, self.out_features, device=x.device, dtype=x.dtype)

        for expert_id, expert in enumerate(self.experts):
            mask = topk_idx == expert_id  # (B, T, k)
            if not mask.any():
                continue
            b_idx, t_idx, k_idx = mask.nonzero(as_tuple=True)
            routed_x = x[b_idx, t_idx]  # (N, D)
            expert_out = expert(routed_x)  # (N, out_features)
            weights = topk_weights[b_idx, t_idx, k_idx].unsqueeze(-1)  # (N, 1)
            output[b_idx, t_idx] += expert_out * weights

        self._aux_loss = self.load_balancing_loss(logits, topk_idx)
        return output

    def load_balancing_loss(self, logits: torch.Tensor, topk_idx: torch.Tensor) -> torch.Tensor:
        """DeepSeekMoE/Switch-style load-balancing loss.

        Args:
            logits: Router logits of shape ``(B, T, E)``.
            topk_idx: Selected expert indices of shape ``(B, T, k)``.

        Returns:
            Scalar load-balancing loss ``E * sum_i importance_i * load_i``.
        """
        b, t, num_experts = logits.shape
        k = topk_idx.shape[-1]

        probs = torch.softmax(logits, dim=-1)  # (B, T, E)
        importance = probs.mean(dim=(0, 1))  # (E,)

        one_hot = F.one_hot(topk_idx, num_classes=num_experts)  # (B, T, k, E)
        load = one_hot.sum(dim=(0, 1, 2)).float() / (b * t * k)  # (E,)

        return num_experts * torch.sum(importance * load)

    def get_aux_loss(self) -> torch.Tensor | float:
        """Return the load-balancing loss from the most recent forward pass."""
        return self._aux_loss


def apply_moe_peft(
    model: nn.Module,
    expert_counts: dict[str, int],
    top_k: int,
    expert_kwargs: dict[str, dict[str, object]],
    controller_noise: bool = False,
) -> nn.Module:
    """Replace every attention ``qkv`` projection in ``model`` with a ``MoEModule``.

    Args:
        model: A timm vision transformer (modified in place).
        expert_counts: Number of experts per family.
        top_k: Number of experts activated per token.
        expert_kwargs: Per-family peft keyword arguments.
        controller_noise: Whether routers add exploration noise in training.

    Returns:
        The same ``model``, with its ``qkv`` layers swapped for routed experts.
    """
    named_modules = dict(model.named_modules())
    for name, module in list(model.named_modules()):
        if not name.endswith("qkv"):
            continue
        if "." in name:
            parent_name, child_name = name.rsplit(".", 1)
            parent = named_modules[parent_name]
        else:
            parent, child_name = model, name

        moe_layer = MoEModule(
            base_linear=module,
            expert_counts=expert_counts,
            top_k=top_k,
            expert_kwargs=expert_kwargs,
            controller_noise=controller_noise,
        )
        setattr(parent, child_name, moe_layer)

    return model