"""DeepSets 编码器 + LSTM 解码器（方案 B 自包含分发包的平铺依赖副本）。"""

import torch
import torch.nn as nn

from vocab import Vocab
import config


class DeepSetsEncoder(nn.Module):
    def __init__(self, input_dim: int = None):
        super().__init__()
        dims = config.ENCODER_MLP_DIMS
        H = config.ENCODER_OUT_DIM

        layers = []
        in_dim = input_dim if input_dim is not None else 0
        for out_dim in dims:
            layers.append(nn.Linear(in_dim, out_dim))
            layers.append(nn.ReLU())
            in_dim = out_dim
        self.phi = nn.Sequential(*layers)

        rho_layers = [
            nn.Linear(H * 2, H),
            nn.ReLU(),
            nn.Linear(H, H),
        ]
        self.rho = nn.Sequential(*rho_layers)

    def forward(self, X, y, n_samples, var_counts=None):
        """
        X: [B, N_max, D_MAX]  (padded, 0 for missing samples and missing vars)
        y: [B, N_max]
        n_samples: [B]  每个 batch 的真实样本数
        var_counts: [B] 每个样本的真实变量维度（未使用，保留兼容签名）
        """
        B, N_max, D = X.shape
        v = torch.cat([X, y.unsqueeze(-1)], dim=-1)

        h = self.phi(v)  # [B, N_max, H]

        # 只对样本维做 mask（padding 的行置零 / 置 -inf），broadcast 到 H 维
        sample_mask = torch.arange(N_max, device=X.device).unsqueeze(0) < n_samples.unsqueeze(1)
        sample_mask = sample_mask.unsqueeze(-1).float()  # [B, N_max, 1]

        h = h * sample_mask

        h_mean = h.sum(dim=1) / n_samples.unsqueeze(-1).float().clamp(min=1)
        h_max = h.masked_fill(sample_mask == 0, float("-inf")).max(dim=1).values

        z = torch.cat([h_mean, h_max], dim=-1)
        z = self.rho(z)
        return z


class ExpressionLSTM(nn.Module):
    def __init__(self, vocab_size: int):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, config.LSTM_EMBEDDING_DIM)
        self.lstm = nn.LSTM(
            input_size=config.LSTM_EMBEDDING_DIM,
            hidden_size=config.LSTM_HIDDEN_DIM,
            num_layers=config.LSTM_NUM_LAYERS,
            batch_first=True,
            dropout=config.DROPOUT if config.LSTM_NUM_LAYERS > 1 else 0,
        )
        self.proj = nn.Linear(config.LSTM_HIDDEN_DIM, vocab_size)

    def forward(self, token_ids, hidden, cell):
        """
        token_ids: [B, T]
        hidden, cell: [1, B, H]
        returns: logits [B, T, V], hidden, cell
        """
        embedded = self.embedding(token_ids)
        output, (hidden, cell) = self.lstm(embedded, (hidden, cell))
        logits = self.proj(output)
        return logits, hidden, cell


class SymbolicRegressionModel(nn.Module):
    def __init__(self, vocab: Vocab):
        super().__init__()
        self.vocab = vocab
        self.encoder = DeepSetsEncoder(input_dim=config.D_MAX + 1)
        self.lstm = ExpressionLSTM(vocab.vocab_size)

    def forward(self, batch):
        X = batch["X"]
        y = batch["y"]
        n_samples = batch["n_samples"]
        token_ids = batch["token_ids"]
        var_counts = batch["var_counts"]

        z = self.encoder(X, y, n_samples, var_counts)

        hidden = z.unsqueeze(0)
        cell = z.unsqueeze(0)

        logits, _, _ = self.lstm(token_ids[:, :-1], hidden, cell)
        return logits

    @torch.no_grad()
    def generate(self, X, y, n_samples, var_counts=None, max_len: int = None):
        self.eval()
        if max_len is None:
            max_len = config.T_MAX

        B = X.size(0)
        if var_counts is None:
            var_counts = torch.full((B,), X.size(-1), device=X.device)

        z = self.encoder(X, y, n_samples, var_counts)

        hidden = z.unsqueeze(0)
        cell = z.unsqueeze(0)

        token = torch.full((B, 1), fill_value=self.vocab.sos_id, device=X.device)
        generated = [token]

        for _ in range(max_len):
            embedded = self.lstm.embedding(token)
            output, (hidden, cell) = self.lstm.lstm(embedded, (hidden, cell))
            logits = self.lstm.proj(output)
            token = logits.argmax(dim=-1)

            generated.append(token)

            if (token == self.vocab.eos_id).all():
                break

        all_tokens = torch.cat(generated, dim=1)
        return all_tokens