"""
=============================================================================
CS230 · 第 16 讲共享模块：beam search、BLEU、CTC 的折叠规则（从零实现）
=============================================================================

对应讲义：lectures/lecture16_word_embeddings_seq2seq_attention.md §9（beam search）、§10（BLEU）、§11（CTC）

-----------------------------------------------------------------------------
【提供的函数】
-----------------------------------------------------------------------------
    beam_search(model, src, B, max_len, alpha)    在 seq2seq_attention.Seq2Seq 上做 beam search；B = 1 就是贪心搜索
    normalized_score(log_p, length, alpha)        长度归一化：(1 / T_y^α) · log P(y | x)
    ngram_counts / modified_precision             修正的 n-gram 精确率（裁剪计数）
    corpus_bleu / sentence_bleu                   BLEU = BP · exp(Σ_n w_n log p_n)
    ctc_collapse                                  CTC 的规则：先合并相邻的重复字符，再去掉 blank
    ctc_brute_force_prob                          把所有对齐方式枚举出来求和，得到 P(y | x)（只能用于很小的例子）

直接运行本文件做自检：BLEU 在 Ng 的 "the cat is on the mat" 例子和 Papineni et al. (2002) 论文的例子上的数字；
CTC 的折叠规则。
    python3 beam_bleu.py          # < 1 秒

【注释标签】【思路】【参数】【返回】【边界】【DL 联系】【Python 小知识】
"""

import itertools
import math
from collections import Counter

import torch

from seq2seq_attention import EOS_ID, SOS_ID, pad_batch


# =============================================================================
# 1. beam search
# =============================================================================
def normalized_score(log_p, length, alpha):
    """
    长度归一化后的分数：(1 / T_y^α) · log P(y | x)。

    【思路】log P(y | x) = Σ_t log P(y<t> | ...) 是很多个负数相加：句子越长，加的负数越多，分数越低。
            所以"不归一化"的目标天然偏爱短句。除以 T_y^α 是一个**启发式**的修正：
            α = 0 不归一化；α = 1 完全归一化（每个词的平均 log 概率）；Ng 说 α = 0.7 是常见的折中。
    【参数】log_p：float；length：T_y（这里**包括** <eos> 这一个词，因为 log_p 里也包括它的概率）；alpha：α
    【返回】float
    """
    return log_p / (length ** alpha)


@torch.no_grad()
def beam_search(model, src, B=3, max_len=15, alpha=0.0):
    """
    beam search：每一步只保留到目前为止 log P 最大的 B 个部分句子。

    【思路】（Ng，C5M3 "Beam Search"）
            第 1 步：算 P(y<1> | x)，留下最好的 B 个词。
            第 t 步：对留下的每个部分句子 y<1..t−1>，算 P(y<t> | x, y<1..t−1>)，一共 B × |V| 个候选，
                    按 log P(y<1..t> | x) = log P(y<1..t−1> | x) + log P(y<t> | ...) 排序，留下最好的 B 个。
            候选里以 <eos> 结尾的句子算"完成了"，放到一边；没完成的继续往下走。全部完成、或者到了 max_len 就停。
            最后在所有完成的句子里，按 normalized_score（长度归一化）挑分数最高的那句。
            · 搜索过程中的剪枝用的是**没有归一化**的 log P（Ng 的讲法也是这样），归一化只用在最后的挑选上。
            · B = 1 时每一步只留一个 —— 这就是贪心搜索。
    【参数】
        model   : seq2seq_attention.Seq2Seq（用它的 encode / decoder_step）
        src     : list[int]，一个源句的编号
        B       : beam 宽度
        max_len : 最多生成几个词（不含 <eos>）
        alpha   : 长度归一化的指数 α
    【返回】(best_tokens, best_logp, finished)
        best_tokens : list[int]，挑出的句子（不含 <eos>）
        best_logp   : float，它的 log P(y | x)（含 <eos> 那一步）
        finished    : list[(tokens, logp)]，所有完成的候选（误差分析和练习用）
    【边界】和广度优先 / 深度优先这类**精确**搜索不同，beam search 不保证找到 argmax_y P(y | x)：
            真正最好的句子可能在某一步排在第 B + 1 名，被剪掉了。B 越大越不容易漏，但越慢（每一步 B × |V| 个候选）。
    【Python 小知识】sorted(..., key=..., reverse=True) 从大到小排；元组 (分数, 下标) 也可以直接比较。
    """
    model.eval()
    x, x_len = pad_batch([src])
    enc1 = model.encode(x, x_len)
    beams = [([], 0.0, enc1["s0"][0])]               # 每个部分句子：(已生成的词, log P, 解码器状态)
    finished = []
    for _ in range(max_len + 1):
        if not beams:
            break
        n = len(beams)
        enc = {k: v.expand(n, *v.shape[1:]) for k, v in enc1.items()}    # 同一句源句，复制 n 份
        y_prev = torch.tensor([b[0][-1] if b[0] else SOS_ID for b in beams], dtype=torch.long)
        s_prev = torch.stack([b[2] for b in beams])
        lp, s_new, _ = model.decoder_step(y_prev, s_prev, enc)          # (n, |V|)
        total = lp + torch.tensor([b[1] for b in beams])[:, None]        # 每个候选的 log P(y<1..t> | x)
        top_vals, top_idx = torch.topk(total.flatten(), min(B, total.numel()))
        V = lp.shape[1]
        new_beams = []
        for val, flat in zip(top_vals.tolist(), top_idx.tolist()):
            i, w = divmod(flat, V)                                        # 第 i 个部分句子 + 第 w 个词
            tokens = beams[i][0] + [w]
            if w == EOS_ID:
                finished.append((tokens[:-1], val))
            elif len(tokens) <= max_len:
                new_beams.append((tokens, val, s_new[i]))
        beams = new_beams
    if not finished:                                                      # 到了 max_len 还没有 <eos>
        finished = [(b[0], b[1]) for b in beams]
    best = max(finished, key=lambda f: normalized_score(f[1], len(f[0]) + 1, alpha))
    return best[0], best[1], finished


# =============================================================================
# 2. BLEU
# =============================================================================
def ngram_counts(tokens, n):
    """
    一句话里每个 n-gram 出现的次数。

    【参数】tokens：词的列表；n：n-gram 的长度
    【返回】Counter：键是长度为 n 的元组，值是次数
    【Python 小知识】Counter 是一个"会数数的字典"；zip(*[tokens[i:] for i in range(n)]) 把错开 0..n−1 位的
            n 个列表并排，得到所有相邻的 n 元组。例如 n = 2：zip(tokens, tokens[1:])。
    """
    return Counter(zip(*[tokens[i:] for i in range(n)]))


def modified_precision(candidate, references, n):
    """
    修正的 n-gram 精确率的分子和分母（Papineni et al. 2002）。

    【思路】普通的精确率 = 候选译文的 n-gram 里有多少出现在参考译文里。它有一个漏洞：
            "the the the the the the the" 的每个词都出现在参考译文里，精确率 7/7 = 100%。
            修正：每个 n-gram 的计数**裁剪**（clip）到它在**任何一个**参考译文里出现的最多次数：
                Count_clip(g) = min( 候选里 g 的次数, max_参考 参考里 g 的次数 )
            参考译文 "the cat is on the mat" 里 the 最多出现 2 次，所以上面那句的修正精确率是 2/7。
    【参数】candidate：词的列表；references：list[词的列表]；n
    【返回】(clipped, total)：裁剪后的匹配数、候选里 n-gram 的总数（corpus_bleu 把各句的分子、分母分别相加）
    """
    cand = ngram_counts(candidate, n)
    max_ref = Counter()
    for ref in references:
        for g, c in ngram_counts(ref, n).items():
            max_ref[g] = max(max_ref[g], c)
    clipped = sum(min(c, max_ref[g]) for g, c in cand.items())
    return clipped, sum(cand.values())


def brevity_penalty(c, r):
    """
    简短惩罚 BP = 1（c > r 时），否则 exp(1 − r / c)。c = 候选译文的总长度，r = 参考译文的"有效"总长度。

    【思路】修正精确率只惩罚"多说"（多余的词匹配不上），不惩罚"少说"：一句只有 "the cat" 的译文，
            精确率可以是 100%。BP 专门惩罚太短的译文。
    【边界】c = 0（什么都没输出）时返回 0。
    """
    if c == 0:
        return 0.0
    return 1.0 if c > r else math.exp(1.0 - r / c)


def corpus_bleu(candidates, references_list, max_n=4):
    """
    语料级 BLEU（Papineni et al. 2002）：BLEU = BP · exp( Σ_{n=1}^{N} (1/N) log p_n )。

    【思路】
        · p_n：把**所有句子**的裁剪计数加起来，除以所有句子的 n-gram 总数（不是每句算一个精确率再平均）。
        · 1–4 gram 的精确率取**几何平均**（log 的平均再 exp）：任何一个 p_n 是 0，BLEU 就是 0。
        · r：每句选和候选长度**最接近**的那个参考译文的长度（一样近时取短的），再加起来 —— 论文叫"best match length"。
    【参数】candidates：list[词的列表]；references_list：list[list[词的列表]]，和 candidates 一一对应；max_n：N
    【返回】dict："bleu"、"p"（[p_1, ..., p_N]）、"bp"、"c"、"r"
    【DL 联系】BLEU 就是机器翻译的"单一数字评估指标"（第 10 讲 §2）：便宜、可以自动算、和人工评价大体一致，
            所以能让 Idea → Code → Experiment 的循环转得快。它不能直接当损失函数来训练（不可导），这和 F1 一样。
    """
    num = [0] * max_n
    den = [0] * max_n
    c_total, r_total = 0, 0
    for cand, refs in zip(candidates, references_list):
        for n in range(1, max_n + 1):
            cl, tot = modified_precision(cand, refs, n)
            num[n - 1] += cl
            den[n - 1] += tot
        c_total += len(cand)
        r_total += min((abs(len(r) - len(cand)), len(r)) for r in refs)[1]
    p = [num[i] / den[i] if den[i] > 0 else 0.0 for i in range(max_n)]
    bp = brevity_penalty(c_total, r_total)
    bleu = bp * math.exp(sum(math.log(pi) for pi in p) / max_n) if min(p) > 0 else 0.0
    return {"bleu": bleu, "p": p, "bp": bp, "c": c_total, "r": r_total}


def sentence_bleu(candidate, references, max_n=4):
    """单句的 BLEU（就是只有一句的 corpus_bleu）。单句的 4-gram 经常是 0，所以单句 BLEU 很粗糙。"""
    return corpus_bleu([candidate], [references], max_n)


# =============================================================================
# 3. CTC
# =============================================================================
def ctc_collapse(seq, blank="_"):
    """
    CTC 的折叠函数 β：先把相邻的重复字符合并成一个，再去掉 blank。

    【思路】语音识别里输入是很多帧（10 秒 × 100 帧/秒 = 1000 帧），输出只有几十个字符。CTC 让网络**每一帧**都输出一个
            字符或 blank（"_"），再用 β 把它折叠成最终的文字（Graves et al. 2006）：
                β("ttt_h_eee___ ___qqq__") = "the q"     （Ng 视频里的例子，中间那个是空格字符）
                β("hh_e_ll_lo") = "hello"                 两个 l 之间必须隔一个 blank，否则会被合并成一个
    【参数】seq：字符串或字符列表；blank：blank 符号
    【返回】折叠后的字符串
    【Python 小知识】itertools.groupby(seq) 把相邻的相同元素分成一组；for ch, _ in groupby(seq) 依次得到每组的那个字符。
    """
    return "".join(ch for ch, _ in itertools.groupby(seq) if ch != blank)


def ctc_brute_force_prob(probs, alphabet, target, blank="_"):
    """
    P(target | x) = Σ_{所有对齐 a，β(a) = target} Π_t p_t(a_t)：把所有可能的逐帧输出枚举一遍。

    【思路】CTC 假设各帧的输出在给定输入时**条件独立**，所以一个对齐（逐帧的字符序列）的概率就是各帧概率的乘积；
            一句文字的概率是所有能折叠成它的对齐的概率之和。
            枚举需要 K^T 次（K 个字符、T 帧）——T = 1000 时完全不可能。真正的 CTC 用动态规划（前向算法），
            计算量只有 O(T · |target|)。这里只在很小的例子上枚举，用来验证 PyTorch 的 ctc_loss 算的就是这个和。
    【参数】probs：(T, K) 每帧在 K 个字符上的概率（每行加起来是 1）；alphabet：长度 K 的字符串（含 blank）；target：目标字符串
    【返回】(prob, n_alignments)：概率和能折叠成 target 的对齐个数
    """
    T, K = len(probs), len(alphabet)
    total, count = 0.0, 0
    for a in itertools.product(range(K), repeat=T):
        if ctc_collapse([alphabet[i] for i in a], blank) == target:
            p = 1.0
            for t, i in enumerate(a):
                p *= float(probs[t][i])
            total += p
            count += 1
    return total, count


def _selftest():
    print("BLEU 自检 —— Ng 的例子（C5M3 \"Bleu Score\"）：")
    refs = ["the cat is on the mat".split(), "there is a cat on the mat".split()]
    cand = "the the the the the the the".split()
    cl, tot = modified_precision(cand, refs, 1)
    print(f"    候选 the the the the the the the：普通精确率 7/7，修正精确率 {cl}/{tot}")
    cand2 = "the cat the cat on the mat".split()
    cl, tot = modified_precision(cand2, refs, 2)
    print(f"    候选 the cat the cat on the mat：修正的 bigram 精确率 {cl}/{tot}")
    print("BLEU 自检 —— Papineni et al. (2002) 的例子：")
    p_refs = ["it is a guide to action that ensures that the military will forever heed party commands".split(),
              "it is the guiding principle which guarantees the military forces always being under the command of the party".split(),
              "it is the practical guide for the army always to heed the directions of the party".split()]
    c1 = "it is a guide to action which ensures that the military always obeys the commands of the party".split()
    c2 = "it is to insure the troops forever hearing the activity guidebook that party direct".split()
    for name, c in (("候选 1", c1), ("候选 2", c2)):
        u = modified_precision(c, p_refs, 1)
        b = modified_precision(c, p_refs, 2)
        print(f"    {name}：unigram {u[0]}/{u[1]}，bigram {b[0]}/{b[1]}")
    print("CTC 折叠规则：")
    for s in ["ttt_h_eee___ ___qqq__", "hh_e_ll_lo", "hh_e_lll_oo", "_h_el_lo_"]:
        print(f"    β({s!r}) = {ctc_collapse(s)!r}")


if __name__ == "__main__":
    _selftest()