"""Set Transformer 集合编码器。
输入/输出形状均为 [B, n_tokens, d_model]，
n_tokens 可变（支持 3 或 4）。
"""

import torch
import torch.nn as nn


class MAB(nn.Module):
    """多头注意力块：MAB(Q, K) = LN(H + FFN(H))，H = LN(Q + MultiHeadAttn(Q, K, K))。"""

    def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1, ln: bool = True) -> None:
        super().__init__()
        assert d_model % num_heads == 0, "d_model 必须能被 num_heads 整除"
        self.attn = nn.MultiheadAttention(
            d_model, num_heads, dropout=dropout, batch_first=True
        )
        self.norm1 = nn.LayerNorm(d_model) if ln else nn.Identity()
        self.norm2 = nn.LayerNorm(d_model) if ln else nn.Identity()
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_model * 4, d_model),
            nn.Dropout(dropout),
        )

    def forward(self, q: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
        """
        Args:
            q: [B, n_q, d_model]
            k: [B, n_k, d_model]

        Returns:
            [B, n_q, d_model]
        """
        attn_out, _ = self.attn(q, k, k)
        h = self.norm1(q + attn_out)
        return self.norm2(h + self.ffn(h))


class SAB(nn.Module):
    """集合自注意力块：SAB(X) = MAB(X, X)。"""

    def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1, ln: bool = True) -> None:
        super().__init__()
        self.mab = MAB(d_model, num_heads, dropout, ln)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """[B, n, d_model] -> [B, n, d_model]"""
        return self.mab(x, x)


class ISAB(nn.Module):
    """诱导点集合注意力块：用 m 个可学习诱导点降低大集合的注意力复杂度。"""

    def __init__(
        self,
        d_model: int,
        num_heads: int,
        num_inducing: int = 16,
        dropout: float = 0.1,
        ln: bool = True,
    ) -> None:
        super().__init__()
        # 可学习诱导点
        self.inducing = nn.Parameter(torch.empty(1, num_inducing, d_model))
        nn.init.xavier_uniform_(self.inducing)
        self.mab_in = MAB(d_model, num_heads, dropout, ln)
        self.mab_out = MAB(d_model, num_heads, dropout, ln)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """[B, n, d_model] -> [B, n, d_model]"""
        inducing = self.inducing.expand(x.size(0), -1, -1)  # [B, m, d_model]
        h = self.mab_in(inducing, x)  # [B, m, d_model]
        return self.mab_out(x, h)     # [B, n, d_model]


class SetTransformer(nn.Module):
    """对化学 token 集合做 N 层集合自注意力，形状保持 [B, n_tokens, d_model]。"""

    def __init__(
        self,
        d_model: int,
        num_heads: int = 8,
        n_layers: int = 4,
        dropout: float = 0.1,
        block: str = "sab",
        num_inducing: int = 16,
        ln: bool = True,
    ) -> None:
        """
        Args:
            d_model: token 维度
            num_heads: 注意力头数，d_head = d_model / num_heads
            n_layers: 集合注意力层数
            dropout: dropout 比例
            block: 注意力块类型，"sab"（全自注意力）或 "isab"（诱导点）
            num_inducing: ISAB 的诱导点数量（block="isab" 时生效）
            ln: 是否使用 LayerNorm
        """
        super().__init__()
        self.block_type = block

        self.layers = nn.ModuleList()
        for _ in range(n_layers):
            if block == "isab":
                self.layers.append(ISAB(d_model, num_heads, num_inducing, dropout, ln))
            else:
                self.layers.append(SAB(d_model, num_heads, dropout, ln))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: [B, n_tokens, d_model] 化学 token 集合

        Returns:
            [B, n_tokens, d_model] 集合编码后的 token
        """
        for layer in self.layers:
            x = layer(x)
        return x