"""
decoding.py —— 解码：从"下一个 token 的分布"到一段文字
=====================================================================

【这个文件要回答的问题】
    语言模型每一步只给出一个分布 p(y_t | y_<t)。要真正"说话"，还需要一个
    **解码算法**决定每一步选哪个 token。第 1 讲 slides 列出了 6 种选项：
      1. 贪心（greedy）：每步选概率最大的那个（argmax）；
      2. 纯采样：按模型给的分布随机抽；
      3. 带温度的采样：先把分布"变尖"或"变平"再抽；
      4. top-k 采样：只在概率最大的 k 个 token 里抽；
      5. top-p（核采样，nucleus sampling）：只在"累计概率刚好达到 p"的那一小撮里抽；
      6. 束搜索（beam search）：同时保留 B 条最有希望的候选，近似寻找整句概率最大的序列。
    另外还有两类"附加旋钮"：频率惩罚 / 存在惩罚（压低已经说过的 token）、停止条件。

【设计】
    所有函数都只依赖一个"打分函数" next_logprobs(history_ids) -> 长度 V 的 log 概率数组，
    所以它们既能接神经网络（neural_lm.NextTokenModel），也能接手写的玩具分布
    （05 号脚本里的束搜索小例子），和真实 LLM 的解码代码结构完全相同。

【LLM 联系】
    OpenAI / Anthropic 等 API 里的 temperature、top_p、top_k、frequency_penalty、
    presence_penalty、max_tokens、stop 参数，就是这里实现的东西。
    （不同厂商支持的参数不完全一样，但含义一致。）
"""

import math

import numpy as np

from lm_math import log_softmax


# =====================================================================
# 1. 改造分布的三个"过滤器"：温度、top-k、top-p
# =====================================================================

def apply_temperature(logprobs, temperature):
    """
    温度：把 log 概率除以 T 再重新归一化，p_T(i) ∝ p(i)^(1/T)。

    【参数】logprobs：长度 V 的 log 概率（或 logits，效果一样）；temperature：T > 0。
    【返回】新的 log 概率数组（已归一化）。
    【直觉】T < 1：大的更大、小的更小，分布变"尖"（更保守）；T -> 0 就是贪心。
            T > 1：分布变"平"（更大胆），T -> ∞ 就是在所有 token 里均匀乱选。
    【边界】T 必须 > 0。log 概率为 -inf 的 token（概率 0）不管 T 是多少都还是 0。
    """
    return log_softmax(logprobs, temperature=temperature)


def top_k_filter(logprobs, k):
    """
    top-k：只保留概率最大的 k 个 token，其余设为 -inf（概率 0），再重新归一化。

    【参数】k：保留的个数（>= 1）。k >= V 时等于不过滤。
    【返回】新的 log 概率数组。
    【边界】若第 k 名和第 k+1 名概率恰好相等，用 argsort 的稳定排序决定，结果是确定的。
    【Python 小知识】np.argsort(-x, kind="stable") 得到"从大到小"的下标顺序。
    """
    k = int(k)
    if k >= len(logprobs):
        return np.asarray(logprobs, dtype=np.float64)
    order = np.argsort(-logprobs, kind="stable")
    out = np.full(len(logprobs), -np.inf)
    keep = order[:k]
    out[keep] = logprobs[keep]
    return log_softmax(out)


def top_p_filter(logprobs, p):
    """
    top-p（核采样）：按概率从大到小排序，保留最少的若干个 token，使它们的累计概率 >= p。

    【参数】p：在 (0, 1] 之间。p = 1 等于不过滤。
    【返回】新的 log 概率数组。
    【为什么不用固定的 k】模型很确定时（比如 "the children" 后面几乎一定是动词），
        候选应该很少；模型不确定时（选哪种水果），候选应该很多。top-p 让候选个数
        随分布的形状**自动调整**，top-k 做不到。
    【边界】至少保留 1 个 token（概率最大的那个），所以 p 再小也不会"无词可选"。
    """
    probs = np.exp(logprobs)
    order = np.argsort(-probs, kind="stable")
    cum = np.cumsum(probs[order])
    # 第一个使累计概率 >= p 的位置（含），再 +1 得到要保留的个数
    # 【Python 小知识】np.searchsorted(有序数组, x) 返回 x 应该插入的位置，即第一个 >= x 的下标。
    n_keep = int(np.searchsorted(cum, p - 1e-12)) + 1
    n_keep = min(max(n_keep, 1), len(probs))
    out = np.full(len(logprobs), -np.inf)
    keep = order[:n_keep]
    out[keep] = logprobs[keep]
    return log_softmax(out)


def apply_penalties(logprobs, counts, frequency_penalty=0.0, presence_penalty=0.0):
    """
    频率惩罚与存在惩罚（OpenAI API 文档里的公式）：
        logit[j] ← logit[j] − c[j] · frequency_penalty − 1[c[j] > 0] · presence_penalty
    其中 c[j] 是 token j 在**已经生成的文字**里出现的次数。

    【参数】
        logprobs：长度 V 的 log 概率；counts：长度 V 的整数数组（每个 token 已出现的次数）。
        frequency_penalty：出现越多次，扣得越多（线性）；
        presence_penalty：只要出现过就扣一次固定的分（不管出现几次）。
    【返回】惩罚并重新归一化之后的 log 概率。
    【边界】惩罚为负数时反而**鼓励**重复（API 允许这样做）。
            注意惩罚对所有 token 一视同仁——包括 "the"、"." 这种本来就该反复出现的词，
            所以惩罚太大会把句子弄得不合语法（05 号脚本会演示）。
    """
    counts = np.asarray(counts)
    adjusted = (np.asarray(logprobs, dtype=np.float64)
                - counts * frequency_penalty
                - (counts > 0) * presence_penalty)
    return log_softmax(adjusted)


# =====================================================================
# 2. 逐 token 生成
# =====================================================================

def generate(next_logprobs, rng, prefix=(), max_new_tokens=30, greedy=False,
             temperature=1.0, top_k=None, top_p=None,
             frequency_penalty=0.0, presence_penalty=0.0,
             stop_ids=(), banned_ids=(), record=None):
    """
    从前缀出发，一个 token 一个 token 地生成，直到遇到停止 token 或达到长度上限。

    【参数】
        next_logprobs：函数，输入历史 id 列表，返回下一个 token 的 log 概率（长度 V）。
        rng：np.random.Generator（显式传入，保证可复现）。
        prefix：已有的 token id（提示 prompt），不计入惩罚计数。
        max_new_tokens：最多生成多少个新 token（停止条件之一）。
        greedy：True 时每步取 argmax（此时 temperature / top_k / top_p 都不起作用）。
        temperature、top_k、top_p：见上面三个过滤器；None 表示不用。处理顺序：
            惩罚 → 温度 → top-k → top-p → 采样（和常见实现一致）。
        frequency_penalty、presence_penalty：见 apply_penalties。
        stop_ids：遇到这些 token 就停（停止 token 本身会保留在输出里）。
        banned_ids：永远不允许生成的 token（例如演示"复读"时禁止 <eos>）。
        record：可选的列表；若提供，每一步会把"这一步模型给出的原始 log 概率"追加进去，方便分析。
    【返回】(新生成的 id 列表, 停止原因)。停止原因是 "stop"（遇到停止 token）或 "length"（达到上限）。
    【边界】greedy=True 时完全确定，不消耗随机数。
    """
    history = list(prefix)
    new = []
    counts = None
    for _ in range(max_new_tokens):
        lp = np.asarray(next_logprobs(history), dtype=np.float64)
        if record is not None:
            record.append(lp)
        if counts is None:
            counts = np.zeros(len(lp), dtype=np.int64)
        if banned_ids:
            lp = lp.copy()
            lp[list(banned_ids)] = -np.inf
            lp = log_softmax(lp)
        if frequency_penalty or presence_penalty:
            lp = apply_penalties(lp, counts, frequency_penalty, presence_penalty)
        if greedy:
            tok = int(np.argmax(lp))
        else:
            if temperature != 1.0:
                lp = apply_temperature(lp, temperature)
            if top_k is not None:
                lp = top_k_filter(lp, top_k)
            if top_p is not None:
                lp = top_p_filter(lp, top_p)
            probs = np.exp(lp)
            probs = probs / probs.sum()             # 消除浮点误差，rng.choice 要求严格加起来等于 1
            tok = int(rng.choice(len(probs), p=probs))
        new.append(tok)
        history.append(tok)
        counts[tok] += 1
        if tok in stop_ids:
            return new, "stop"
    return new, "length"


# =====================================================================
# 3. 束搜索
# =====================================================================

def beam_search(next_logprobs, prefix=(), beam_size=4, max_new_tokens=30, stop_ids=(),
                length_alpha=0.0, return_all=False):
    """
    束搜索：每一步只保留得分最高的 beam_size 条部分序列（"束"），近似寻找整句 log 概率最大的序列。

    【思路】
        第 t 步：对每条还没结束的束，尝试接上词表里的每一个 token，得分 = 原得分 + log p(token)；
        已经结束（以停止 token 结尾）的束原样保留；从所有候选里挑出得分最高的 beam_size 条。
        直到所有束都结束，或者达到长度上限。
        beam_size = 1 就退化成贪心解码；beam_size = V^T 就是穷举（不可行）。
    【参数】
        next_logprobs、prefix、max_new_tokens、stop_ids：同 generate。
        beam_size：束宽 B（>= 1）。
        length_alpha：长度归一化指数 α。最终排序用 得分 / (长度^α)。α = 0 就是原始的总 log 概率
            （偏爱短句：每多一个词总得分就再减一点）；α = 1 是"平均每个 token 的 log 概率"。
        return_all：True 时返回最终全部束（按排序分数从高到低），否则只返回最好的一条。
    【返回】return_all=False：(id 列表, 总 log 概率)；return_all=True：[(id 列表, 总 log 概率, 排序分数), ...]。
    【边界】
        - 为了简单，这里在每一步都只保留 B 条（不管是否结束），这是最常见的教科书写法。
        - 达到长度上限还没结束的束也会参与最终排序。
    """
    beams = [(0.0, [], False)]          # (总 log 概率, 新生成的 id, 是否已结束)
    for _ in range(max_new_tokens):
        candidates = []
        for score, ids, done in beams:
            if done:
                candidates.append((score, ids, True))
                continue
            lp = np.asarray(next_logprobs(list(prefix) + ids), dtype=np.float64)
            for tok in np.flatnonzero(np.isfinite(lp)):      # 只考虑概率 > 0 的 token
                tok = int(tok)
                candidates.append((score + lp[tok], ids + [tok], tok in stop_ids))
        # 按总得分从高到低排序；得分相同按 id 序列排序，保证结果确定
        candidates.sort(key=lambda c: (-c[0], c[1]))
        beams = candidates[:beam_size]
        if all(done for _, _, done in beams):
            break

    def rank_score(score, ids):
        return score / (max(len(ids), 1) ** length_alpha)

    ranked = sorted(beams, key=lambda b: (-rank_score(b[0], b[1]), b[1]))
    if return_all:
        return [(ids, score, rank_score(score, ids)) for score, ids, _ in ranked]
    best = ranked[0]
    return best[1], best[0]


def sequence_logprob(next_logprobs, ids, prefix=()):
    """
    用链式法则计算一段 token 在模型下的总 log 概率：Σ_t log p(ids[t] | prefix + ids[:t])。

    【返回】float（可能是 -inf）。
    """
    total = 0.0
    history = list(prefix)
    for tok in ids:
        total += float(next_logprobs(history)[tok])
        history.append(tok)
        if math.isinf(total):
            break
    return total