"""
MoE (Mixture-of-Experts) 模块：sample-level + 跨模态路由。

设计要点（与现有 8-token 架构对齐）：
    - Router 输入：tab token 池化后的 [B, d]（即配方/实验条件向量）。
    - Expert 输入：chem token flatten 后的 [B, T_chem * d]（化学侧 token）。
    - 路由粒度：每个样本只算一次 router → gates [B, K]。
    - Top-k 稀疏激活 + 可选训练态 jitter 噪声。
    - 返回 load-balancing aux loss，由 trainer 端按权重加进总 loss。

输出:
    F_moe: [B, d_model]   与单个 token 同维度，便于追加到 fusion 序列中。
    extras: dict           诊断与监控信息（aux loss、gates 等）。
"""

from typing import Dict, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F


class MoEAttentionPool(nn.Module):
    """与 FusionLayer 同款的 attention pooling，参数独立。

    将一组 token [B, T, d] 池化成单个查询向量 [B, d]，用作 router 的输入。
    """

    def __init__(self, d_model: int) -> None:
        super().__init__()
        self.d_model = d_model
        self.query = nn.Parameter(torch.randn(1, 1, d_model))
        self.proj = nn.Linear(d_model, d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: [B, T, d_model]

        Returns:
            [B, d_model]
        """
        B = x.size(0)
        q = self.query.expand(B, -1, -1)                              # [B, 1, d]
        k = self.proj(x)                                              # [B, T, d]
        scores = torch.bmm(q, k.transpose(1, 2)) / (self.d_model ** 0.5)
        weights = F.softmax(scores, dim=-1)                           # [B, 1, T]
        return torch.bmm(weights, x).squeeze(1)                       # [B, d]


class MoERouter(nn.Module):
    """单层 softmax router + Top-k 稀疏激活。

    Args:
        d_model: router 输入维度。
        n_experts: 专家数量 K。
        top_k: 每个样本激活的专家数。
        jitter_noise: 训练态加在 logits 上的均匀噪声幅度，用于探索；评估态自动关闭。
    """

    def __init__(
        self,
        d_model: int,
        n_experts: int,
        top_k: int = 2,
        jitter_noise: float = 0.0,
    ) -> None:
        super().__init__()
        if not 1 <= top_k <= n_experts:
            raise ValueError(f"top_k 必须在 [1, n_experts={n_experts}]，收到 {top_k}")
        self.n_experts = n_experts
        self.top_k = top_k
        self.jitter_noise = jitter_noise
        self.linear = nn.Linear(d_model, n_experts)

    def forward(
        self, q: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Args:
            q: [B, d_model] 路由查询向量。

        Returns:
            gates:       [B, n_experts] top-k 后再归一化的稀疏概率。
            probs_full:  [B, n_experts] 未掩码的 softmax 概率（用于 aux loss）。
            expert_mask: [B, n_experts] 0/1 掩码，标记被激活的专家。
        """
        logits = self.linear(q)                                       # [B, K]
        if self.training and self.jitter_noise > 0.0:
            noise = (torch.rand_like(logits) - 0.5) * self.jitter_noise
            logits = logits + noise

        probs_full = F.softmax(logits, dim=-1)                        # [B, K]

        topk_vals, topk_idx = probs_full.topk(self.top_k, dim=-1)     # [B, k]
        expert_mask = torch.zeros_like(probs_full)
        expert_mask.scatter_(1, topk_idx, 1.0)

        gates = probs_full * expert_mask
        gates = gates / gates.sum(dim=-1, keepdim=True).clamp(min=1e-9)
        return gates, probs_full, expert_mask


class MoEExpert(nn.Module):
    """单个专家 MLP：in_dim → hidden_dim → out_dim。"""

    def __init__(
        self, in_dim: int, hidden_dim: int, out_dim: int, dropout: float = 0.1,
    ) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, out_dim),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class MoEBlock(nn.Module):
    """
    Sample-level 跨模态 MoE。

    流程:
        1. tab tokens [B, T_tab, d] -> attention pool -> q_tab [B, d]
        2. q_tab -> Router -> gates [B, K] (+ probs_full / expert_mask)
        3. chem tokens [B, T_chem, d] -> flatten -> [B, T_chem * d]
        4. K 个 Expert MLP 并行处理 -> stack [B, K, d]
        5. 加权求和 -> F_moe [B, d]
        6. 计算 load-balancing aux loss

    Args:
        d_model: token 维度。
        n_chem_tokens: chem 侧 token 数（决定 expert 输入维度）。
        n_experts: 专家数量 K。
        top_k: 每个样本激活的专家数。
        expert_hidden_mult: expert 中间层维度 = expert_hidden_mult * d_model。
        dropout: expert 内部 dropout。
        jitter_noise: router 训练态噪声幅度，0.0 表示关闭。
    """

    def __init__(
        self,
        d_model: int,
        n_chem_tokens: int = 4,
        n_experts: int = 4,
        top_k: int = 2,
        expert_hidden_mult: int = 2,
        dropout: float = 0.1,
        jitter_noise: float = 0.0,
    ) -> None:
        super().__init__()
        self.d_model = d_model
        self.n_chem_tokens = n_chem_tokens
        self.n_experts = n_experts
        self.top_k = top_k

        self.tab_pool = MoEAttentionPool(d_model)
        self.router = MoERouter(
            d_model, n_experts, top_k=top_k, jitter_noise=jitter_noise,
        )

        expert_in = d_model * n_chem_tokens
        expert_hidden = d_model * expert_hidden_mult
        self.experts = nn.ModuleList([
            MoEExpert(expert_in, expert_hidden, d_model, dropout=dropout)
            for _ in range(n_experts)
        ])

    def forward(
        self,
        chem: torch.Tensor,
        tab: torch.Tensor,
    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        """
        Args:
            chem: [B, T_chem, d_model] 化学侧 token（router 不看）。
            tab:  [B, T_tab,  d_model] 配方/实验侧 token（router 看）。

        Returns:
            F_moe: [B, d_model]
            extras: {
                "lb_loss": 标量 load-balancing aux loss（带梯度），
                "gates":   [B, K] 稀疏归一后的门控（detached），
                "probs":   [B, K] 原始 softmax 概率（detached），
            }
        """
        if chem.size(1) != self.n_chem_tokens:
            raise ValueError(
                f"chem token 数不匹配：期望 {self.n_chem_tokens}, "
                f"实际 {chem.size(1)}"
            )

        q_tab = self.tab_pool(tab)                                     # [B, d]
        gates, probs_full, expert_mask = self.router(q_tab)            # 各 [B, K]

        flat = chem.flatten(start_dim=1)                               # [B, T_chem * d]
        expert_outs = torch.stack(
            [expert(flat) for expert in self.experts], dim=1,
        )                                                              # [B, K, d]

        F_moe = (gates.unsqueeze(-1) * expert_outs).sum(dim=1)         # [B, d]

        lb_loss = self._load_balancing_loss(probs_full, expert_mask)

        return F_moe, {
            "lb_loss": lb_loss,
            "gates": gates.detach(),
            "probs": probs_full.detach(),
        }

    def _load_balancing_loss(
        self,
        probs_full: torch.Tensor,
        expert_mask: torch.Tensor,
    ) -> torch.Tensor:
        """Switch Transformer 风格的 load-balancing loss。

            f_i = 该 batch 中 expert_i 被激活的样本占比（top-k mask 的均值）
            p_i = 该 batch 中 expert_i 的 softmax 概率均值
            loss = K * Σ_i f_i * p_i

        理想情况下 f_i 与 p_i 都接近 1/K，loss ≈ 1。
        """
        f = expert_mask.mean(dim=0)                                   # [K]
        p = probs_full.mean(dim=0)                                    # [K]
        return self.n_experts * (f * p).sum()