"""
num_lab.py —— 数字在计算机里的样子：整数、各种浮点格式、舍入、求和、量化（前置讲 P2 的工具箱）
=====================================================================

【这个文件要回答的问题】
    LLM 系统里到处是"换一种数字格式"：训练用 BF16、累加用 FP32、推理用 FP8 甚至 INT4。
    要理解这些选择，得先能回答：一种格式能表示多大、多小、多精细的数？一个数存进去会变成什么？
    这个模块不依赖任何深度学习库，用几十行公式实现了**任意**"符号 + 指数 + 尾数"浮点格式的
    编码、解码和舍入，并在 02 号脚本里和 PyTorch 的真实实现逐个核对。

【本文件提供】
    1. 整数      : int_range、twos_complement_bits、wrap_int
    2. 浮点格式  : FloatFormat、FORMATS（FP32、TF32、FP16、BF16、FP8 E4M3、FP8 E5M2、FP4 E2M1）、
                   format_stats、round_to（按"就近舍入、平局取偶"存进某种格式）、decode、fields、all_values
    3. 求和      : sum_sequential、sum_pairwise、sum_kahan（都可以指定"每一步都按某种格式舍入"）
    4. 量化预告  : absmax_quantize、zeropoint_quantize、dequantize

【怎样读这里的公式】（以 FP16 为例：1 位符号 s、5 位指数 E、10 位尾数 M）
    规格化数（E 不是全 0 也不是全 1）：值 = (−1)^s × (1 + M / 2^10) × 2^(E − 15)，15 叫"偏置"（bias）
    次正规数（E 全 0）              ：值 = (−1)^s × (0 + M / 2^10) × 2^(1 − 15)
    E 全 1                           ：无穷大（M = 0）或 NaN（M ≠ 0）——FP8 E4M3 和 FP4 例外（见 FORMATS）

【系统联系】
    每种格式的"指数位 / 尾数位"分配决定了它的范围和精度：BF16 的范围和 FP32 一样大、精度只有 3 位十进制，
    所以训练很少溢出；FP16 精度更高但范围小，要配合损失缩放。量化（第 19–20 讲）、混合精度训练（第 11、18 讲）、
    FP8 训练与推理（第 21 讲 FlashAttention-3）都建立在这些事实上。
"""

import math
from dataclasses import dataclass

import numpy as np


# =====================================================================
# 1. 整数
# =====================================================================

def int_range(bits, signed=True):
    """
    bits 位整数能表示的范围。

    【返回】(最小值, 最大值)：有符号（补码）是 [−2^(bits−1), 2^(bits−1) − 1]，无符号是 [0, 2^bits − 1]
    【例子】int8 → (−128, 127)；uint8 → (0, 255)；int4 → (−8, 7)
    """
    if signed:
        return -(1 << (bits - 1)), (1 << (bits - 1)) - 1
    return 0, (1 << bits) - 1


def twos_complement_bits(x, bits):
    """
    把整数 x 写成 bits 位的补码比特串。

    【思路】补码：非负数照常写；负数 x 写成 2^bits + x 的二进制。于是 −1 是全 1，最高位就是"符号位"。
        补码的好处是加法电路不用区分正负：直接按无符号相加，丢掉溢出的进位，结果就对。
    【边界】x 超出 int_range(bits) 时抛出 ValueError。
    """
    lo, hi = int_range(bits, True)
    if not lo <= x <= hi:
        raise ValueError(f"{x} 超出 {bits} 位有符号整数的范围 [{lo}, {hi}]")
    return format(x & ((1 << bits) - 1), f"0{bits}b")


def wrap_int(x, bits, signed=True):
    """
    模拟固定位数整数的"回绕"（wrap-around）：只保留最低 bits 位，再按补码 / 无符号解释。

    【例子】wrap_int(127 + 1, 8) = −128；wrap_int(255 + 1, 8, signed=False) = 0
    【系统联系】C、CUDA、NumPy 的整数溢出都是这样悄悄回绕的，不会报错。
    """
    x &= (1 << bits) - 1
    if signed and x >= (1 << (bits - 1)):
        x -= 1 << bits
    return x


# =====================================================================
# 2. 浮点格式
# =====================================================================

@dataclass(frozen=True)
class FloatFormat:
    """
    一种"符号 + 指数 + 尾数"的二进制浮点格式。

    【属性】
        name     : 名字
        exp_bits : 指数位数
        man_bits : 尾数位数（不含隐含的那个 1）
        special  : 指数全 1 时怎样解释——
                   "ieee" : 尾数为 0 是 ±inf，否则是 NaN（FP32 / FP16 / BF16 / FP8 E5M2）
                   "fn"   : 没有 inf，只有"指数全 1 且尾数全 1"是 NaN，其余都是普通的数（FP8 E4M3fn）
                   "none" : 没有 inf 也没有 NaN，所有比特模式都是普通的数（FP4 E2M1）
        torch_dtype_name : PyTorch 里对应的 dtype 名（用来核对）；None 表示 PyTorch 没有独立的这种张量类型
        note     : 一句话用途
    """
    name: str
    exp_bits: int
    man_bits: int
    special: str = "ieee"
    torch_dtype_name: str = None
    note: str = ""

    @property
    def bits(self):
        """总位数 = 1（符号）+ 指数位 + 尾数位。"""
        return 1 + self.exp_bits + self.man_bits

    @property
    def bias(self):
        """指数偏置 = 2^(指数位 − 1) − 1。例如 FP32 是 127，FP16 是 15。"""
        return (1 << (self.exp_bits - 1)) - 1


FP32 = FloatFormat("FP32", 8, 23, "ieee", "float32", "传统的单精度；优化器状态、累加器")
TF32 = FloatFormat("TF32", 8, 10, "ieee", None, "张量核心内部用：FP32 的范围 + FP16 的精度（存储时仍占 32 位）")
FP16 = FloatFormat("FP16", 5, 10, "ieee", "float16", "半精度；范围小，训练要配合损失缩放")
BF16 = FloatFormat("BF16", 8, 7, "ieee", "bfloat16", "\"脑浮点\"：FP32 的范围、很低的精度；当今训练的主流")
FP8_E4M3 = FloatFormat("FP8 E4M3", 4, 3, "fn", "float8_e4m3fn", "前向的权重和激活")
FP8_E5M2 = FloatFormat("FP8 E5M2", 5, 2, "ieee", "float8_e5m2", "反向的梯度（需要更大的范围）")
FP4_E2M1 = FloatFormat("FP4 E2M1", 2, 1, "none", None, "Blackwell 的 FP4（通常配合每块一个缩放因子）")
FORMATS = [FP32, TF32, FP16, BF16, FP8_E4M3, FP8_E5M2, FP4_E2M1]


def max_finite(fmt):
    """
    最大的有限值。

    【思路】IEEE 式：指数全 1 被 inf / NaN 占用，最大值用"指数全 1 减 1、尾数全 1"。
        fn 式（E4M3fn）：指数全 1 也是普通数，只有尾数全 1 是 NaN，所以最大值的尾数是"全 1 减 1"。
        none 式（E2M1）：指数和尾数都取全 1。
    """
    m = fmt.man_bits
    e_all_ones = (1 << fmt.exp_bits) - 1
    if fmt.special == "ieee":
        return (2 - 2.0 ** -m) * 2.0 ** (e_all_ones - 1 - fmt.bias)
    if fmt.special == "fn":
        return (1 + (2 ** m - 2) / 2 ** m) * 2.0 ** (e_all_ones - fmt.bias)
    return (2 - 2.0 ** -m) * 2.0 ** (e_all_ones - fmt.bias)


def format_stats(fmt):
    """
    一种格式的关键数字。

    【返回】字典：
        max         : 最大有限值
        min_normal  : 最小的正规格化数 2^(1 − bias)
        min_sub     : 最小的正次正规数 2^(1 − bias − 尾数位)
        eps         : 机器精度 2^(−尾数位)——1 和"比 1 大的下一个数"之间的距离
        digits      : 约有几位十进制有效数字 ≈ (尾数位 + 1) × log10(2)
    """
    return {"max": max_finite(fmt), "min_normal": 2.0 ** (1 - fmt.bias),
            "min_sub": 2.0 ** (1 - fmt.bias - fmt.man_bits), "eps": 2.0 ** -fmt.man_bits,
            "digits": (fmt.man_bits + 1) * math.log10(2)}


def round_to(x, fmt):
    """
    把一个实数 x 存进格式 fmt：就近舍入，恰好在两个数正中间时取尾数为偶数的那个（round half to even）。

    【思路】
        1. 找到 x 所在的"二进位区间" [2^E, 2^(E+1))，在这个区间里格式能表示的数是等间距的，
           间距 q = 2^(E − 尾数位)。比最小规格化数还小时，E 固定为最小指数（次正规数区间的间距不再变小）。
        2. 把 x 除以 q、按 Python 的 round（本身就是"平局取偶"）取整，再乘回 q。
        3. 超出最大值：IEEE 式格式变成 ±inf；fn / none 式格式饱和到 ±最大值
           （这是本机 PyTorch 2.14 把 FP32 转成 float8_e4m3fn 时的行为；不同库的约定不同，有的会变成 NaN）。
    【参数】x：Python 浮点数；fmt：FloatFormat
    【返回】舍入后的值（Python 浮点数，可以精确表示 fmt 里的任何数）
    【边界】NaN 原样返回；±0 保留符号；inf 在 IEEE 式格式里原样返回，在其他格式里饱和。
    【Python 小知识】math.frexp(a) 把正数 a 写成 m × 2^e，0.5 ≤ m < 1。
    """
    if x != x:
        return x
    if x == 0 or math.isinf(x):
        if math.isinf(x) and fmt.special != "ieee":
            return math.copysign(max_finite(fmt), x)
        return x
    sign = -1.0 if x < 0 else 1.0
    a = abs(x)
    _, e = math.frexp(a)
    E = max(e - 1, 1 - fmt.bias)                 # 规格化时 a ∈ [2^E, 2^(E+1))；太小就用最小指数（次正规）
    q = 2.0 ** (E - fmt.man_bits)                # 这个区间里相邻两个可表示数的间距
    r = round(a / q) * q
    mx = max_finite(fmt)
    if r > mx:
        if fmt.special == "ieee":
            return sign * math.inf
        return sign * mx
    return sign * r


def fields(x, fmt):
    """
    把 x 存进 fmt 后的三段比特：(符号位, 指数字段, 尾数字段, "s eeeee mmmmmmmmmm" 形式的字符串)。

    【边界】NaN 用"指数全 1、尾数全 1"表示；inf 用"指数全 1、尾数 0"表示（只有 IEEE 式格式有 inf）。
    """
    v = round_to(x, fmt)
    s = 1 if (v < 0 or (v == 0 and math.copysign(1.0, v) < 0)) else 0
    m_bits, e_bits = fmt.man_bits, fmt.exp_bits
    if v != v:
        E, M = (1 << e_bits) - 1, (1 << m_bits) - 1
    elif math.isinf(v):
        E, M = (1 << e_bits) - 1, 0
    elif v == 0:
        E, M = 0, 0
    else:
        a = abs(v)
        _, e = math.frexp(a)
        if e - 1 < 1 - fmt.bias:                 # 次正规数
            E = 0
            M = int(round(a / 2.0 ** (1 - fmt.bias - m_bits)))
        else:
            E = e - 1 + fmt.bias
            M = int(round((a / 2.0 ** (e - 1) - 1) * (1 << m_bits)))
    s_str = f"{s} {format(E, f'0{e_bits}b')} {format(M, f'0{m_bits}b')}"
    return s, E, M, s_str


def decode(s, E, M, fmt):
    """
    比特 → 数值（fields 的逆运算），按格式的 special 规则处理指数全 1 的情况。

    【返回】Python 浮点数（可能是 inf 或 NaN）
    """
    e_all = (1 << fmt.exp_bits) - 1
    m_all = (1 << fmt.man_bits) - 1
    sign = -1.0 if s else 1.0
    if E == e_all:
        if fmt.special == "ieee":
            return sign * math.inf if M == 0 else math.nan
        if fmt.special == "fn" and M == m_all:
            return math.nan
    if E == 0:
        return sign * (M / (1 << fmt.man_bits)) * 2.0 ** (1 - fmt.bias)
    return sign * (1 + M / (1 << fmt.man_bits)) * 2.0 ** (E - fmt.bias)


def all_values(fmt):
    """
    枚举一种小格式（≤ 8 位）的全部比特模式，返回所有有限值（含 ±0）组成的列表。

    【边界】超过 8 位的格式会抛出 ValueError（FP16 有 65536 个模式，用不着全部列出）。
    """
    if fmt.bits > 8:
        raise ValueError("只枚举 8 位以内的格式")
    vals = []
    for s in (0, 1):
        for E in range(1 << fmt.exp_bits):
            for M in range(1 << fmt.man_bits):
                v = decode(s, E, M, fmt)
                if not (math.isinf(v) or v != v):
                    vals.append(v)
    return vals


# =====================================================================
# 3. 求和：同样的数，换一种加法顺序
# =====================================================================

def _rounder(fmt):
    """返回"把一个数舍入到 fmt"的函数；fmt 为 None 时不舍入（用 Python 的 double）。"""
    if fmt is None:
        return lambda v: v
    return lambda v: round_to(v, fmt)


def sum_sequential(values, fmt=None):
    """
    从左到右逐个相加，**每一步**的部分和都舍入到 fmt（模拟"累加器就是这种格式"）。

    【思路】IEEE 754 规定：两个浮点数相加的结果，等于先精确相加、再舍入一次。这里先用 double 精确相加
        （对 16 位以内的格式、正常大小的数，double 足够精确），再用 round_to 舍入——就模拟了一个 fmt 格式的加法器。
    【参数】values：可迭代的数（会先各自舍入到 fmt）；fmt：FloatFormat 或 None（None = 用 double）
    """
    r = _rounder(fmt)
    s = 0.0
    for v in values:
        s = r(s + r(v))
    return s


def sum_pairwise(values, fmt=None):
    """
    成对（树形）求和：先两两相加，再把结果两两相加……直到只剩一个。每一步都舍入到 fmt。

    【为什么更准】每次相加的两个数大小相近，舍入误差的累积从 O(n) 降到 O(log n) 量级。
        GPU 上的并行归约天然就是这个形状（第 3 讲）。
    """
    r = _rounder(fmt)
    level = [r(v) for v in values]
    if not level:
        return 0.0
    while len(level) > 1:
        nxt = [r(level[i] + level[i + 1]) for i in range(0, len(level) - 1, 2)]
        if len(level) % 2:
            nxt.append(level[-1])
        level = nxt
    return level[0]


def sum_kahan(values, fmt=None):
    """
    Kahan 补偿求和：多用一个变量 c 记住"每一步被舍入掉的那一点"，下一步再加回去。

    【思路】y = x − c；t = s + y（这里会丢掉 y 的低位）；c = (t − s) − y（恰好算出丢掉的部分，取负）；s = t。
        所有运算都在 fmt 里做，却能得到接近"高一倍精度"的结果。
    """
    r = _rounder(fmt)
    s, c = 0.0, 0.0
    for v in values:
        y = r(r(v) - c)
        t = r(s + y)
        c = r(r(t - s) - y)
        s = t
    return s


# =====================================================================
# 4. 量化预告：把浮点数映射成小整数
# =====================================================================

def absmax_quantize(x, bits=8):
    """
    对称（absmax）量化：缩放因子 scale = max|x| / (2^(bits−1) − 1)，q = round(x / scale)。

    【参数】x：NumPy 数组；bits：整数位数
    【返回】(q 整数数组, scale)
    【边界】x 全为 0 时 scale 取 1，避免除以 0。
    【系统联系】第 19 讲的第一个量化方法；一个离群值就会把 scale 撑大，让其余的数都挤在很少几个整数上。
    """
    qmax = (1 << (bits - 1)) - 1
    amax = float(np.max(np.abs(x)))
    scale = amax / qmax if amax > 0 else 1.0
    q = np.clip(np.round(x / scale), -qmax, qmax).astype(np.int32)
    return q, scale


def zeropoint_quantize(x, bits=8):
    """
    非对称（零点）量化：把 [min, max] 线性映射到 [0, 2^bits − 1]，q = round(x / scale) + zero_point。

    【返回】(q 整数数组, scale, zero_point)
    【为什么要零点】ReLU 之后的激活全是非负数，对称量化会浪费一半的整数（负数那一半）。
    """
    qmax = (1 << bits) - 1
    lo, hi = float(np.min(x)), float(np.max(x))
    scale = (hi - lo) / qmax if hi > lo else 1.0
    zp = int(round(-lo / scale))
    q = np.clip(np.round(x / scale) + zp, 0, qmax).astype(np.int32)
    return q, scale, zp


def dequantize(q, scale, zero_point=0):
    """反量化：x ≈ (q − zero_point) × scale。"""
    return (q.astype(np.float64) - zero_point) * scale