"""
dlm.py —— 第 6 讲的两个小语言模型：掩码扩散语言模型（MDLM）与自回归语言模型（AR）
=====================================================================================

【这个模块提供什么】
    - TinyTransformer：一个很小的 Transformer（默认 3 层、宽 64、4 个头，约 15 万参数）。
      两个语言模型用**完全相同**的网络，只有两处不同：
          掩码扩散：双向注意力（每个位置看得到所有位置），输入是部分被 [MASK] 的句子，输出每个位置的后验；
          自回归  ：因果注意力（只看得到左边），输入是 [BOS] + 句子去掉最后一个词，输出"下一个词"的分布。
    - train_mdlm（原课讲义 Algorithm 8 的掩码版本）、train_ar（标准的下一个词预测）。
    - 评估：ar_nll（精确的负对数似然）、mdlm_nll_fn（交给 ctmc_lib.elbo_by_mask_count 算 ELBO）。
    - 采样：sample_ar（从左到右 d 次前向）；掩码扩散的采样用 ctmc_lib 里的 euler_sample / ordered_unmask_sample，
      这里的 mdlm_posterior_fn 把网络包装成它们需要的 posterior_fn。

【流与扩散联系】掩码扩散的网络**不吃时间 t**。原因（讲义 §4.2）：掩码路径下，后验 p_{1|t}(z^i | x)
    只取决于"哪些位置已揭开、揭开的是什么"，和 t 无关——t 的信息已经包含在 [MASK] 的个数里了。
    MDLM（Sahoo et al. 2024）、RADD（Ou et al. 2024）都指出了这一点；LLaDA 等大模型也不输入 t。
"""

import time

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

from ctmc_lib import DT, MixturePath
from toy_language import BOS_ID, MASK_ID, SEQ_LEN, V


class Block(nn.Module):
    """
    一个 pre-LN 的 Transformer 块：x + Attn(LN(x))，再 x + MLP(LN(x))。

    【参数】dim：宽度；heads：注意力头数（dim 要能被 heads 整除）
    【Python 小知识】F.scaled_dot_product_attention(q, k, v, is_causal=True) 是 PyTorch 2.0 起的
        注意力函数；is_causal=True 时位置 i 只看得到位置 ≤ i（自回归需要的"不偷看未来"）。
    """

    def __init__(self, dim, heads):
        super().__init__()
        self.heads = heads
        self.ln1 = nn.LayerNorm(dim)
        self.qkv = nn.Linear(dim, 3 * dim)
        self.proj = nn.Linear(dim, dim)
        self.ln2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))

    def forward(self, x, causal):
        n, L, D = x.shape
        q, k, v = self.qkv(self.ln1(x)).view(n, L, 3, self.heads, D // self.heads).permute(2, 0, 3, 1, 4)
        a = F.scaled_dot_product_attention(q, k, v, is_causal=causal)          # (n, heads, L, D/heads)
        x = x + self.proj(a.transpose(1, 2).reshape(n, L, D))
        return x + self.mlp(self.ln2(x))


class TinyTransformer(nn.Module):
    """
    小 Transformer：token 嵌入 + 位置嵌入 → depth 个块 → LayerNorm → 线性层输出 n_out 个 logits。

    【参数】
        n_in    : 输入词表大小（20 个真 token + [MASK] + [BOS] = 22）
        n_out   : 输出类别数（只输出 20 个真 token；模型从不预测 [MASK] / [BOS]）
        seq_len : 序列长度（位置嵌入的个数）
        causal  : True = 自回归（因果注意力）；False = 掩码扩散（双向注意力）
    【返回】forward(x)：x 是 (n, L) long，返回 (n, L, n_out) 的 logits
    """

    def __init__(self, n_in=V + 2, n_out=V, seq_len=SEQ_LEN, dim=64, depth=3, heads=4, causal=False):
        super().__init__()
        self.causal = causal
        self.tok = nn.Embedding(n_in, dim)
        self.pos = nn.Parameter(torch.randn(seq_len, dim) * 0.02)
        self.blocks = nn.ModuleList([Block(dim, heads) for _ in range(depth)])
        self.ln = nn.LayerNorm(dim)
        self.out = nn.Linear(dim, n_out)

    def forward(self, x):
        h = self.tok(x) + self.pos[: x.shape[1]]
        for b in self.blocks:
            h = b(h, self.causal)
        return self.out(self.ln(h))


def count_params(model):
    return sum(p.numel() for p in model.parameters())


# =====================================================================
# 训练
# =====================================================================

def train_mdlm(model, lang, steps=3000, batch=256, lr=1e-3, seed=0, weighting="none", log_every=None):
    """
    训练掩码扩散语言模型（原课讲义 Algorithm 8，p_init = δ_[MASK]，κ_t = t）。

    每一步：
        z ~ p_data（玩具语言的句子）；t ~ U[0, 1)；x ~ p_t(· | z)（每个位置以概率 1 − t 换成 [MASK]）；
        损失 = 被掩码位置上的交叉熵 −log p^θ(z^i | x) 的加权平均。

    【参数】
        weighting : "none" —— 每个被掩码的 token 权重相同（实践中常用；最优解仍是真实后验）；
                    "elbo" —— 按 ELBO 的权重 κ̇/(1−κ) = 1/(1−t) 加权（训练目标就是负 ELBO）
        log_every : 每隔多少步记录平均损失
    【返回】(logs: [{step, loss}], elapsed 秒)
    【边界】
        - 未被掩码的位置不计入损失：它们的后验就是自己，没有可学的（Algorithm 8 写的是对所有位置求和，
          两者只差一个常数）。
        - t 接近 1 时几乎没有 [MASK]，一个批里可能出现整句都没有 [MASK] 的情况，它对损失没有贡献。
        - "elbo" 加权时 1/(1−t) 在 t → 1 发散；这里把 t 截断到 0.999。
    """
    g = torch.Generator().manual_seed(seed)
    path = MixturePath(V, "mask", mask_id=MASK_ID)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    logs, running = [], []
    t0 = time.perf_counter()
    model.train()
    for step in range(1, steps + 1):
        z = lang.sample(batch, g)
        t = torch.rand(batch, 1, generator=g, dtype=DT).clamp(max=0.999)
        x = path.sample_conditional(z, t, g)
        masked = x == MASK_ID
        ce = F.cross_entropy(model(x).transpose(1, 2), z, reduction="none")      # (batch, L)
        if weighting == "elbo":
            w = (1 / (1 - t)).float().expand_as(ce)
            loss = (ce * masked * w).sum() / (batch * SEQ_LEN)
        else:
            loss = (ce * masked).sum() / masked.sum().clamp(min=1)
        opt.zero_grad()
        loss.backward()
        opt.step()
        running.append(loss.item())
        if log_every and step % log_every == 0:
            logs.append({"step": step, "loss": sum(running) / len(running)})
            running = []
    model.eval()
    return logs, time.perf_counter() - t0


def ar_inputs(z):
    """自回归的输入：[BOS] + z 去掉最后一个 token（"老师强制" teacher forcing）。"""
    bos = torch.full((z.shape[0], 1), BOS_ID, dtype=torch.long)
    return torch.cat([bos, z[:, :-1]], dim=1)


def train_ar(model, lang, steps=3000, batch=256, lr=1e-3, seed=0, log_every=None):
    """
    训练自回归语言模型：损失 = Σ_i −log p^θ(z^i | z^{<i}) 的平均（每个 token）。

    【参数】同 train_mdlm
    【返回】(logs, elapsed)
    【边界】PAD 也要预测（模型要学会"句子在哪里结束"），和掩码扩散一致。
    """
    g = torch.Generator().manual_seed(seed)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    logs, running = [], []
    t0 = time.perf_counter()
    model.train()
    for step in range(1, steps + 1):
        z = lang.sample(batch, g)
        loss = F.cross_entropy(model(ar_inputs(z)).transpose(1, 2), z)
        opt.zero_grad()
        loss.backward()
        opt.step()
        running.append(loss.item())
        if log_every and step % log_every == 0:
            logs.append({"step": step, "loss": sum(running) / len(running)})
            running = []
    model.eval()
    return logs, time.perf_counter() - t0


# =====================================================================
# 评估与采样
# =====================================================================

@torch.no_grad()
def ar_nll(model, z):
    """自回归模型的精确负对数似然 −log p^θ(z) = Σ_i −log p^θ(z^i | z^{<i})（nats / 句），(n,) float64。"""
    logp = F.log_softmax(model(ar_inputs(z)).double(), dim=2)
    return -logp.gather(2, z.unsqueeze(2)).squeeze(2).sum(dim=1)


def mdlm_nll_fn(model, z):
    """
    把掩码扩散模型包装成 ctmc_lib.elbo_by_mask_count 需要的 nll_fn：x → (n, d) 的 −log p^θ(z^i | x)。
    """
    @torch.no_grad()
    def fn(x):
        logp = F.log_softmax(model(x).double(), dim=2)
        return -logp.gather(2, z.unsqueeze(2)).squeeze(2)
    return fn


def mdlm_posterior_fn(model):
    """
    把掩码扩散模型包装成采样器需要的 posterior_fn：(x, t) → (n, d, V) 概率。t 不用（见模块说明）。
    """
    @torch.no_grad()
    def fn(x, t):
        return F.softmax(model(x).double(), dim=2)
    return fn


@torch.no_grad()
def sample_ar(model, n, generator=None):
    """
    从自回归模型从左到右采样：每生成一个 token 做一次前向（共 d 次）。

    【返回】(x: (n, d) long, nfe: 前向次数 = d)
    【边界】为了代码简单，这里每次都把整个前缀重新算一遍；真实系统用 KV 缓存，
        每一步只算新 token，总计算量与一次完整的前向相当（讲义 §5.2）。
    """
    z = torch.zeros(n, SEQ_LEN, dtype=torch.long)
    for i in range(SEQ_LEN):
        logits = model(ar_inputs(z))[:, i]
        z[:, i] = torch.multinomial(F.softmax(logits.double(), dim=1), 1, generator=generator).squeeze(1)
    return z, SEQ_LEN