"""
=============================================================================
CS230 · 第 16 讲共享模块：编码器-解码器（seq2seq），有 / 没有 Bahdanau 注意力（PyTorch）
=============================================================================

对应讲义：lectures/lecture16_word_embeddings_seq2seq_attention.md §7、§8（代码 04、05 都用它）

-----------------------------------------------------------------------------
【结构】
-----------------------------------------------------------------------------
    编码器 Encoder：嵌入 → 双向 GRU。每个输入位置 t' 得到一个"注解"（annotation）
                    a<t'> = [a→<t'>; a←<t'>]，形状 (B, T_x, 2H)：前向读到 t' 为止 + 反向读到 t' 为止。
    解码器 Decoder：每一步 t 输入"上一个输出词 y<t−1>"和一个上下文向量 c<t>，GRU 更新状态 s<t>，再预测 y<t>：
                    s<t> = GRU([emb(y<t−1>); c<t>], s<t−1>)
                    P(y<t> | x, y<1..t−1>) = softmax(W_o [s<t>; c<t>] + b_o)
        · 没有注意力（Cho et al. 2014 的写法）：c<t> 每一步都是**同一个**向量 —— 编码器读完整句话后的总结
          c = [a→<T_x>; a←<1>]。整句话的信息都要挤进这一个固定长度的向量。
        · 有注意力（Bahdanau, Cho & Bengio 2015）：每一步重新算一个 c<t>：
                e<t,t'> = vᵀ tanh(W_s s<t−1> + W_a a<t'>)          一个小网络打分："现在该看第 t' 个位置吗？"
                α<t,t'> = softmax_{t'}(e<t,t'>)                    对所有输入位置归一化，Σ_{t'} α<t,t'> = 1
                c<t>    = Σ_{t'} α<t,t'> a<t'>                      注解的加权平均
    两种解码器的其余部分完全一样（同样的 GRU、同样的输出层），只差 c<t> 怎么来，比较才公平。

    Coursera 作业里的注意力小网络是 "Dense(10, tanh) → Dense(1, relu)"，输入是 [s<t−1>; a<t'>] 的拼接；
    这里用 Bahdanau 论文的写法 vᵀ tanh(W_s s + W_a a)，两者是同一个想法。

-----------------------------------------------------------------------------
【提供的函数】
-----------------------------------------------------------------------------
    Seq2Seq                 模型（attention=True / False）
    pad_batch               把长短不一的编号列表补齐成张量
    train_seq2seq           teacher forcing 训练若干个 epoch
    greedy_decode           贪心解码（每一步取概率最大的词），同时返回注意力权重
    sequence_log_prob       给定 x 和一个完整的 y，算 log P(y | x) = Σ_t log P(y<t> | x, y<1..t−1>)
    beam search 在第 05 个脚本的 beam_bleu.py 里，用的是这里的 encode / decoder_step

【形状约定】PyTorch 习惯，每一行一个样本：x (B, T_x)，y (B, T_y)；编号 0 是 <pad>。
【可复现】CPU、单线程（调用方 torch.set_num_threads(1)）、建模型前 torch.manual_seed、打乱顺序用自己的 torch.Generator。
【注释标签】【思路】【参数】【返回】【边界】【DL 联系】【PyTorch 小知识】【Python 小知识】
"""

import math

import torch
from torch import nn

PAD_ID, SOS_ID, EOS_ID = 0, 1, 2         # 和 seq_worlds.Vocab 的约定一致


def pad_batch(seqs, pad_id=PAD_ID):
    """
    把若干个长短不一的编号列表补齐成一个 (B, T_max) 的长整数张量，并返回每个序列的真实长度。

    【参数】seqs：list[list[int]]
    【返回】(tensor (B, T_max), lengths (B,) 长整数张量)
    """
    T = max(len(s) for s in seqs)
    out = torch.full((len(seqs), T), pad_id, dtype=torch.long)
    for i, s in enumerate(seqs):
        out[i, :len(s)] = torch.tensor(s, dtype=torch.long)
    return out, torch.tensor([len(s) for s in seqs], dtype=torch.long)


class Seq2Seq(nn.Module):
    """
    编码器-解码器，attention=True 时用 Bahdanau 注意力，否则每一步用同一个总结向量。

    【参数】
        n_in, n_out   : 输入 / 输出词表大小
        emb           : 嵌入维数
        hidden        : 编码器每个方向的 GRU 单元数 H（注解 a<t'> 是 2H 维）；解码器状态 s<t> 是 2H 维
        attn_dim      : 注意力小网络的隐藏层宽度
        attention     : True / False
    【PyTorch 小知识】nn.GRU(..., bidirectional=True, batch_first=True) 的输出形状是 (B, T, 2H)，
            最后的隐藏状态 h_n 是 (2, B, H)：h_n[0] 是前向读完的状态，h_n[1] 是反向读完（读到第 1 个字符）的状态。
            nn.utils.rnn.pack_padded_sequence 让 GRU 跳过补齐的 <pad>，h_n 才是每个序列"真正读完"时的状态。
    """

    def __init__(self, n_in, n_out, emb=32, hidden=64, attn_dim=64, attention=True):
        super().__init__()
        self.attention = attention
        self.hidden = hidden
        d_ann = 2 * hidden                                           # 注解 a<t'> 的维数
        self.src_emb = nn.Embedding(n_in, emb, padding_idx=PAD_ID)
        self.encoder = nn.GRU(emb, hidden, batch_first=True, bidirectional=True)
        self.init_state = nn.Linear(d_ann, d_ann)                    # s<0> = tanh(W [a→<T_x>; a←<1>])
        self.tgt_emb = nn.Embedding(n_out, emb, padding_idx=PAD_ID)
        self.cell = nn.GRUCell(emb + d_ann, d_ann)                   # 解码器：输入 [emb(y<t−1>); c<t>]
        self.out = nn.Linear(2 * d_ann, n_out)                       # 输出层：[s<t>; c<t>] → 词表上的 logits
        if attention:
            self.W_a = nn.Linear(d_ann, attn_dim, bias=False)        # 作用在注解 a<t'> 上（每句话只算一次）
            self.W_s = nn.Linear(d_ann, attn_dim)                    # 作用在解码器上一步的状态 s<t−1> 上
            self.v = nn.Linear(attn_dim, 1, bias=False)              # vᵀ tanh(...) → 一个分数 e<t,t'>

    # ------------------------------------------------------------------ 编码
    def encode(self, x, x_len):
        """
        读完输入句子。

        【参数】x：(B, T_x) 长整数；x_len：(B,) 真实长度
        【返回】dict：
            "ann"     : (B, T_x, 2H) 每个位置的注解 a<t'>（<pad> 位置是 0）
            "summary" : (B, 2H) 总结向量 [a→<T_x>; a←<1>]（没有注意力的解码器每一步都用它）
            "mask"    : (B, T_x) 布尔，True 表示真实字符（不是 <pad>）
            "keys"    : (B, T_x, attn_dim) W_a a<t'>，预先算好（有注意力时）
            "s0"      : (B, 2H) 解码器的初始状态
        """
        e = self.src_emb(x)
        packed = nn.utils.rnn.pack_padded_sequence(e, x_len, batch_first=True, enforce_sorted=False)
        out, h_n = self.encoder(packed)
        ann, _ = nn.utils.rnn.pad_packed_sequence(out, batch_first=True, total_length=x.shape[1])
        summary = torch.cat([h_n[0], h_n[1]], dim=1)
        enc = {"ann": ann, "summary": summary, "mask": x != PAD_ID, "s0": torch.tanh(self.init_state(summary))}
        if self.attention:
            enc["keys"] = self.W_a(ann)
        return enc

    # ------------------------------------------------------------------ 注意力
    def attend(self, s_prev, enc):
        """
        Bahdanau 注意力：根据解码器上一步的状态 s<t−1>，给每个输入位置打分、归一化、加权平均。

        【参数】s_prev：(B, 2H)；enc：encode 的输出
        【返回】(c, alpha)：上下文向量 (B, 2H) 和注意力权重 (B, T_x)（每一行加起来是 1，<pad> 位置是 0）
        【边界】<pad> 位置的分数设成一个很大的负数（而不是 −inf），softmax 之后是 0，又不会出现 nan。
        【DL 联系】这是本讲最重要的几行：e<t,t'> = vᵀ tanh(W_s s<t−1> + W_a a<t'>)；α = softmax(e)；c = Σ α a。
                计算量：每个输出位置都要和每个输入位置打一次分，一共 T_x × T_y 个 α（讲义 §8.4）。
        """
        scores = self.v(torch.tanh(enc["keys"] + self.W_s(s_prev)[:, None, :])).squeeze(-1)   # (B, T_x)
        scores = scores.masked_fill(~enc["mask"], -1e9)
        alpha = torch.softmax(scores, dim=1)
        c = torch.bmm(alpha[:, None, :], enc["ann"]).squeeze(1)                               # (B, 2H)
        return c, alpha

    # ------------------------------------------------------------------ 解码一步
    def decoder_step(self, y_prev, s_prev, enc):
        """
        解码器走一步。

        【参数】y_prev：(B,) 上一个输出词的编号（第一步是 <sos>）；s_prev：(B, 2H)；enc：encode 的输出
        【返回】(log_probs (B, n_out), s (B, 2H), alpha (B, T_x) 或 None)
        【PyTorch 小知识】torch.log_softmax 直接算 log 概率（内部用 log-sum-exp，第 9 讲 §3.7），比 log(softmax) 稳定。
        """
        if self.attention:
            c, alpha = self.attend(s_prev, enc)
        else:
            c, alpha = enc["summary"], None
        s = self.cell(torch.cat([self.tgt_emb(y_prev), c], dim=1), s_prev)
        logits = self.out(torch.cat([s, c], dim=1))
        return torch.log_softmax(logits, dim=1), s, alpha

    def forward(self, x, x_len, y_in):
        """
        teacher forcing：解码器每一步的输入是**正确答案**的上一个词（而不是模型自己上一步的预测）。

        【参数】x, x_len：输入；y_in：(B, T_y) 解码器的输入 = [<sos>, y<1>, ..., y<T_y−1>]
        【返回】(B, T_y, n_out) 每一步的 log 概率
        """
        enc = self.encode(x, x_len)
        s = enc["s0"]
        outs = []
        for t in range(y_in.shape[1]):
            lp, s, _ = self.decoder_step(y_in[:, t], s, enc)
            outs.append(lp)
        return torch.stack(outs, dim=1)


def _shift_right(y):
    """y = [y1, ..., yT]（含 <eos>）→ 解码器输入 [<sos>, y1, ..., y_{T−1}]。"""
    return torch.cat([torch.full((y.shape[0], 1), SOS_ID, dtype=torch.long), y[:, :-1]], dim=1)


def train_seq2seq(model, src, tgt, epochs=10, batch_size=64, lr=3e-3, seed=0, log_every=0, clip=1.0):
    """
    用 teacher forcing + 交叉熵训练 seq2seq。

    【思路】损失 = 每个目标位置上 −log P(正确的词)，对所有非 <pad> 位置取平均（"条件语言模型"的最大似然）。
            用 Adam；梯度裁剪（clip，第 15 讲讲 RNN 时会讲为什么）防止偶尔的梯度爆炸。
    【参数】
        model    : Seq2Seq
        src, tgt : list[list[int]]，tgt 已经包含末尾的 <eos>
        epochs, batch_size, lr, seed : 训练设置（seed 决定打乱顺序）
        log_every: >0 时每隔这么多个 epoch 打印一次平均损失
        clip     : 梯度范数的上限
    【返回】list[float]：每个 epoch 的平均损失（按目标词数加权）
    【边界】发散检测：损失出现 inf / nan 时立刻抛出异常。
    【PyTorch 小知识】nn.NLLLoss 接收 log 概率（我们的模型已经做了 log_softmax）；ignore_index=PAD_ID 让 <pad> 位置不计入损失。
    """
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    loss_fn = nn.NLLLoss(ignore_index=PAD_ID, reduction="sum")
    gen = torch.Generator().manual_seed(seed)
    history = []
    n = len(src)
    for ep in range(1, epochs + 1):
        model.train()
        perm = torch.randperm(n, generator=gen).tolist()
        total, count = 0.0, 0
        for i in range(0, n, batch_size):
            idx = perm[i:i + batch_size]
            x, x_len = pad_batch([src[j] for j in idx])
            y, _ = pad_batch([tgt[j] for j in idx])
            logp = model(x, x_len, _shift_right(y))
            n_tok = int((y != PAD_ID).sum().item())
            loss = loss_fn(logp.reshape(-1, logp.shape[-1]), y.reshape(-1))
            opt.zero_grad()
            (loss / n_tok).backward()
            nn.utils.clip_grad_norm_(model.parameters(), clip)
            opt.step()
            total += loss.item()
            count += n_tok
        avg = total / count
        if not math.isfinite(avg):
            raise FloatingPointError("训练发散了（损失是 inf / nan）")
        history.append(avg)
        if log_every and ep % log_every == 0:
            print(f"    epoch {ep:>3}   每个目标词的平均损失 {avg:.4f}")
    return history


@torch.no_grad()
def greedy_decode(model, src, max_len=12, batch_size=256):
    """
    贪心解码：每一步取概率最大的词，直到 <eos> 或 max_len。

    【参数】src：list[list[int]]；max_len：最多输出几个词
    【返回】(preds, alphas)
        preds  : list[list[int]]，每句的输出编号（不含 <eos>）
        alphas : list[tensor (T_out, T_x)]，每句每一步的注意力权重（没有注意力时为 None）
    """
    model.eval()
    preds, alphas = [], []
    for i in range(0, len(src), batch_size):
        chunk = src[i:i + batch_size]
        x, x_len = pad_batch(chunk)
        enc = model.encode(x, x_len)
        s = enc["s0"]
        y_prev = torch.full((len(chunk),), SOS_ID, dtype=torch.long)
        steps, att = [], []
        for _ in range(max_len):
            lp, s, alpha = model.decoder_step(y_prev, s, enc)
            y_prev = lp.argmax(dim=1)
            steps.append(y_prev)
            if alpha is not None:
                att.append(alpha)
        steps = torch.stack(steps, dim=1)                     # (B, max_len)
        att = torch.stack(att, dim=1) if att else None        # (B, max_len, T_x)
        for b in range(len(chunk)):
            seq = steps[b].tolist()
            cut = seq.index(EOS_ID) if EOS_ID in seq else len(seq)
            preds.append(seq[:cut])
            alphas.append(att[b, :cut + 1, :len(chunk[b])].clone() if att is not None else None)
    return preds, alphas


@torch.no_grad()
def sequence_log_prob(model, src, tgt, batch_size=256):
    """
    log P(y | x) = Σ_t log P(y<t> | x, y<1..t−1>)：给定一个完整的输出（含 <eos>），模型认为它有多可能。

    【思路】和训练时一样用 teacher forcing 把 y 喂进解码器，把每一步"正确词"的 log 概率加起来。
            beam search 的误差分析（讲义 §9.4）就是比较 log P(y* | x) 和 log P(ŷ | x)。
    【参数】src, tgt：list[list[int]]，tgt 要包含末尾的 <eos>
    【返回】list[float]：每句的 log P(y | x)
    """
    model.eval()
    out = []
    for i in range(0, len(src), batch_size):
        x, x_len = pad_batch(src[i:i + batch_size])
        y, _ = pad_batch(tgt[i:i + batch_size])
        logp = model(x, x_len, _shift_right(y))                                 # (B, T_y, V)
        picked = logp.gather(2, y.unsqueeze(-1)).squeeze(-1)                    # 每一步正确词的 log 概率
        picked = picked.masked_fill(y == PAD_ID, 0.0)
        out += picked.sum(dim=1).tolist()
    return out