"""
lm_math.py —— 语言模型输出层的数学：softmax、log-softmax、交叉熵、困惑度、熵
=====================================================================

【这个文件要回答的问题】
    神经语言模型的最后几步永远是同一套数学：
        logits（每个词一个分数） --softmax--> 概率 --取 -log--> 损失（负对数似然）
    这个文件用 NumPy 从零实现这些函数，并且特别注意**数值稳定性**：
    直接按公式 exp(z) / Σ exp(z) 计算，z 稍大一点（比如 1000）就会溢出成 inf，
    最后得到 nan。所有真实框架（PyTorch 的 F.softmax、F.cross_entropy）内部都
    用了下面这些"减去最大值"的技巧。

【怎样使用】
    from lm_math import softmax, log_softmax, nll, perplexity, entropy
    p = softmax(np.array([2.0, 1.0, 0.1]))            # -> [0.659, 0.242, 0.099]
    p = softmax(logits, temperature=0.5)              # 温度越低越"尖"

【LLM 联系】
    这些函数就是 LLM 输出层的全部数学。第 1 讲的"logits → softmax → 损失 → 解码"
    四块积木，代码都在这里；第 19 讲的 FlashAttention 会再用到 logsumexp 的
    "在线"版本（边读数据边更新最大值）。
"""

import math

import numpy as np


def logsumexp(z, axis=-1):
    """
    数值稳定地计算 log(Σ_i exp(z_i))。

    【思路】先减去最大值 m：log Σ exp(z_i) = m + log Σ exp(z_i - m)。
            减完之后最大的一项是 exp(0) = 1，其余都 <= 1，不会溢出；
            而且至少有一项等于 1，所以求和结果 >= 1，取 log 也不会遇到 log(0)。
    【参数】z：NumPy 数组（任意形状）；axis：沿哪一维求和（默认最后一维 = 词表那一维）。
    【返回】比 z 少一维的数组（axis 那一维被求和掉）；z 是一维时返回一个标量。
    【边界】z 里有 -inf（"这个词绝对不可能"）是允许的：exp(-inf) = 0。
            但如果整行都是 -inf，结果是 -inf（没有任何可能的词）。
    【Python 小知识】np.max(z, axis=axis, keepdims=True) 保留被求最大值的那一维（长度变成 1），
            这样 z - m 时能按"广播"规则自动对齐形状。
    """
    z = np.asarray(z, dtype=np.float64)
    m = np.max(z, axis=axis, keepdims=True)
    m = np.where(np.isfinite(m), m, 0.0)          # 整行 -inf 时 m 也是 -inf，换成 0 避免 (-inf) - (-inf) = nan
    out = m + np.log(np.sum(np.exp(z - m), axis=axis, keepdims=True))
    return np.squeeze(out, axis=axis)


def log_softmax(z, temperature=1.0, axis=-1):
    """
    数值稳定的 log-softmax：log p_i = z_i / T - logsumexp(z / T)。

    【参数】z：logits 数组；temperature：温度 T（> 0）；axis：词表所在的维度。
    【返回】和 z 形状相同的数组，每个元素 <= 0，exp 之后沿 axis 加起来等于 1。
    【边界】T 必须 > 0。T -> 0 的极限是 argmax（见 decoding.py 的贪心解码），这里不处理 T = 0。
    【LLM 联系】训练时我们几乎从不直接算 softmax 再取 log（会先下溢成 0 再 log 出 -inf），
               而是一步算出 log-softmax，PyTorch 的 F.cross_entropy 也是这样做的。
    """
    if temperature <= 0:
        raise ValueError("temperature 必须大于 0；T = 0 请改用 argmax（贪心）")
    z = np.asarray(z, dtype=np.float64) / temperature
    return z - np.expand_dims(logsumexp(z, axis=axis), axis)


def softmax(z, temperature=1.0, axis=-1):
    """
    数值稳定的 softmax：p_i = exp(z_i / T) / Σ_j exp(z_j / T)。

    【参数】同 log_softmax。
    【返回】概率数组，形状同 z，每个元素在 [0, 1]，沿 axis 加起来等于 1。
    【思路】exp(log_softmax)。等价于"减去最大值再 exp 再归一化"。
    【边界】
        - 给所有 logits 加同一个常数，结果不变（平移不变性），所以 logits 本身的"绝对大小"没有意义，
          只有差值有意义。
        - 极端的 logits 差（比如 1000）会让小概率下溢成精确的 0.0，这在数值上是正确的。
    """
    return np.exp(log_softmax(z, temperature=temperature, axis=axis))


def naive_softmax(z):
    """
    **故意写错**的 softmax：直接按课本公式 exp(z) / sum(exp(z))。

    只用于演示数值溢出：z 里有 1000 时 exp(1000) = inf，inf / inf = nan。
    【返回】可能包含 nan 的数组（NumPy 会给出 RuntimeWarning，我们在调用处关掉警告）。
    """
    e = np.exp(np.asarray(z, dtype=np.float64))
    return e / e.sum()


def nll(logits, target, temperature=1.0):
    """
    负对数似然（negative log-likelihood）：-log p(target)。

    【参数】
        logits：一维数组（一个位置的 logits），或二维数组 (位置数, 词表大小)。
        target：正确答案的下标（int），或下标数组（每个位置一个）。
    【返回】一维输入时返回一个 float；二维输入时返回每个位置的 NLL 组成的数组。
    【边界】如果模型给正确答案的概率恰好是 0（logit = -inf），NLL = +inf——
            这就是 n-gram 不平滑时"测试集困惑度无穷大"的原因。
    【LLM 联系】把它在所有位置上取平均，就是 LLM 的训练损失（交叉熵损失）。
    """
    lp = log_softmax(logits, temperature=temperature, axis=-1)
    if lp.ndim == 1:
        return float(-lp[target])
    target = np.asarray(target)
    # 【Python 小知识】lp[np.arange(n), target] 是"花式索引"：第 i 行取第 target[i] 列。
    return -lp[np.arange(lp.shape[0]), target]


def perplexity(total_nll, n_tokens):
    """
    困惑度 = exp(平均每个 token 的 NLL)。

    【参数】total_nll：所有 token 的 NLL 之和（自然对数）；n_tokens：token 个数（> 0）。
    【返回】float，>= 1。total_nll 为 inf 时返回 inf。
    【直觉】困惑度 = "模型平均每一步在多少个等可能的选项里犹豫"。
            一个对 V 个词完全均匀猜测的模型，困惑度恰好是 V。
    """
    if n_tokens <= 0:
        raise ValueError("n_tokens 必须大于 0")
    if math.isinf(total_nll):
        return math.inf
    return math.exp(total_nll / n_tokens)


def entropy(p, axis=-1):
    """
    分布的熵 H(p) = -Σ p_i log p_i（自然对数，单位 nat）。

    【参数】p：概率数组（沿 axis 加起来等于 1）。
    【返回】float 或数组。
    【边界】p_i = 0 的项按约定 0·log 0 = 0（用 np.where 避开 log(0) 的警告）。
    【直觉】熵越大，分布越"平"（越难猜）；exp(熵) 就是这个分布的"有效选项个数"，
            和困惑度是同一个概念。
    """
    p = np.asarray(p, dtype=np.float64)
    safe = np.where(p > 0, p, 1.0)                    # p = 0 的位置随便填个 1，log(1) = 0
    return -np.sum(np.where(p > 0, p * np.log(safe), 0.0), axis=axis)