"""
=============================================================================
CS349E · 前置讲 P3 共享模块：numfmt.py —— 逐位编码 / 解码各种数值格式，模拟低精度运算
=============================================================================

【这个模块要解决什么问题】
    推理硬件的很多"倍数"都来自数值格式：FP8 的算力是 BF16 的 2 倍，NVFP4 又是 FP8 的 2–3 倍。
    但"换一个格式"到底丢掉了什么？要回答它，必须能把一个数**真的**变成某种格式、再变回来，看误差。
    我们的电脑（和 numpy）只原生支持 FP64 / FP32 / FP16，所以这里用纯 numpy 写一个**通用的**浮点格式：

      1. 整数：补码的编码 / 解码、回绕（wrap）和饱和（saturate）两种溢出处理。
      2. FloatFormat：给定"指数几位、尾数几位、偏置多少、有没有 Inf/NaN"，就能
         - round()：把任意实数舍入到这个格式能表示的最近的数（支持 4 种舍入方式）；
         - encode() / decode()：在实数和比特之间来回转换；
         - 列出所有能表示的数（小格式）。
         预置了 FP32、TF32、FP16、BF16、FP8（E4M3 / E5M2）、FP6（E3M2 / E2M3）、FP4（E2M1）。
      3. 块缩放（block scaling）：MX 格式（OCP 规范：每 32 个数共享一个 2 的幂缩放因子）和 NVFP4
         （每 16 个数一个 FP8 缩放因子 + 整个张量一个 FP32 缩放因子）。
      4. 整数量化：仿射（scale + zero point）、对称 / 非对称、按张量 / 按通道 / 按组。
      5. 低精度累加：每做一次加法就舍入一次，模拟"累加器只有 p 位"的硬件；顺序、两两（pairwise）、
         Kahan 补偿求和、以及"每 N 个元素提升到 FP32"（DeepSeek-V3 的做法）。
      6. 数值稳定的 softmax：朴素、减最大值、在线（online）softmax 与它的合并运算；Welford 方差。

【怎样验证它是对的】
    自检（python3 numfmt.py）和 02 脚本会把这里的 FP16 / FP32 编码和 numpy 的硬件实现**逐位**比较，
    BF16 和经典的"加 0x7FFF 再右移 16 位"比较，FP8 和 PyTorch 的 float8 类型比较（装了 torch 才比）。

【依据】IEEE 754-2019；OCP *8-bit Floating Point Specification (OFP8)* Rev 1.0；
    OCP *Microscaling Formats (MX) Specification* v1.0（FP6 / FP4 / E8M0 的定义与转换算法）；
    NVIDIA 技术博客 *Introducing NVFP4*（2025-06）。

【运行】python3 numfmt.py    # 自检
【依赖】numpy。
"""

from dataclasses import dataclass

import numpy as np


# ---------------------------------------------------------------------------
# 1. 整数与补码
# ---------------------------------------------------------------------------
def int_range(bits, signed=True):
    """
    返回 bits 位整数能表示的 (最小值, 最大值)。

    【例子】int_range(8) = (-128, 127)；int_range(8, signed=False) = (0, 255)。
    【为什么有符号的负数多一个】补码把最高位当作 −2^(bits−1)，其余位是正的 2^k，
        所以负数一侧到 −128，正数一侧只到 127（0 占用了一个"非负"的编码）。
    """
    if signed:
        return -(1 << (bits - 1)), (1 << (bits - 1)) - 1
    return 0, (1 << bits) - 1


def twos_complement_encode(v, bits):
    """
    把 Python 整数 v 编成 bits 位的补码，返回无符号的编码（0 … 2^bits − 1）。

    【思路】补码的定义：编码 = v mod 2^bits。对负数，这等价于"取反加一"：−1 → 全 1。
    【边界】v 超出 int_range(bits) 时抛出 ValueError（编码本身不检查溢出，那是运算的事）。
    【Python 小知识】Python 的整数没有位数限制；v & mask 对负数也按"无限多个 1 的补码"计算，
        所以 (-1) & 0xFF == 255，正好是我们要的编码。
    """
    lo, hi = int_range(bits)
    if not lo <= v <= hi:
        raise ValueError(f"{v} 超出 {bits} 位有符号整数的范围 [{lo}, {hi}]")
    return v & ((1 << bits) - 1)


def twos_complement_decode(code, bits):
    """把 bits 位的补码编码还原成整数：最高位为 1 时减去 2^bits。"""
    code &= (1 << bits) - 1
    return code - (1 << bits) if code >> (bits - 1) else code


def bits_str(code, bits, groups=None):
    """
    把编码写成二进制字符串；groups 给出分段的位数，例如 FP32 用 (1, 8, 23) 写成 "0 01111111 000…"。

    【边界】groups 的和必须等于 bits，否则抛出 ValueError。
    """
    s = format(int(code) & ((1 << bits) - 1), f"0{bits}b")
    if groups is None:
        return s
    if sum(groups) != bits:
        raise ValueError("groups 的位数之和必须等于 bits")
    out, i = [], 0
    for g in groups:
        out.append(s[i:i + g])
        i += g
    return " ".join(out)


def wrap_int(v, bits):
    """
    回绕（wraparound）：把整数（或 numpy 整数数组）按 2^bits 取模，落回有符号范围。
    这是 C 语言 / numpy 的 int8、int16 在溢出时的行为（127 + 1 → −128）。
    """
    lo, _ = int_range(bits)
    v = np.asarray(v, dtype=np.int64)
    return (v - lo) % (1 << bits) + lo


def saturate_int(v, bits):
    """饱和（saturation）：超出范围的值截到最大 / 最小值（127 + 1 → 127）。量化里一般用这一种。"""
    lo, hi = int_range(bits)
    return np.clip(np.asarray(v, dtype=np.int64), lo, hi)


# ---------------------------------------------------------------------------
# 2. 通用浮点格式
# ---------------------------------------------------------------------------
ROUNDING_MODES = ("nearest_even", "toward_zero", "nearest_away", "stochastic")


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

    【一个数的值】设符号位 S、指数字段 E（无符号整数）、尾数字段 M：
        - 正规数（E ≥ 1）：(−1)^S × 2^(E − bias) × (1 + M / 2^m)      ← 隐含的"1."
        - 次正规数（E = 0）：(−1)^S × 2^(1 − bias) × (0 + M / 2^m)      ← 没有隐含的 1，填补 0 附近的空隙
        - 特殊值取决于 special：
            "ieee"：E 全 1 留给 Inf（M = 0）和 NaN（M ≠ 0）——FP32 / FP16 / BF16 / FP8 E5M2；
            "fn"  ：没有 Inf，只有 E 全 1 且 M 全 1 是 NaN（"finite + NaN"）——FP8 E4M3，多出 7 个数，最大值 448；
            "none"：没有 Inf 也没有 NaN，所有编码都是数——FP6、FP4（OCP MX 规范）。
    【参数】name：名字；exp_bits / man_bits：指数 / 尾数的位数；bias：指数偏置；special：见上。
    【常用的派生量】见下面的 property：bits、emax、emin、max_value、min_normal、min_subnormal、eps。
    【边界】只支持 bits ≤ 32 的格式（编码用 int64 存）；all_values() 只对 bits ≤ 16 的格式开放。
    【Python 小知识】@dataclass(frozen=True) 自动生成 __init__ 等方法，并且让对象不可修改（可以当字典的键）；
        @property 让方法像属性一样用：FP16.max_value 而不是 FP16.max_value()。
    """

    name: str
    exp_bits: int
    man_bits: int
    bias: int
    special: str = "ieee"

    # ---- 派生量 ----
    @property
    def bits(self):
        """总位数 = 1 + 指数位 + 尾数位。"""
        return 1 + self.exp_bits + self.man_bits

    @property
    def emax(self):
        """正规数的最大（无偏）指数。ieee 格式的 E 全 1 被特殊值占用，所以少 1。"""
        top = (1 << self.exp_bits) - 1
        return top - 1 - self.bias if self.special == "ieee" else top - self.bias

    @property
    def emin(self):
        """正规数的最小（无偏）指数 = 1 − bias；次正规数也用这个指数。"""
        return 1 - self.bias

    @property
    def max_value(self):
        """最大的有限数。fn 格式（E4M3）的 E 全 1、M 全 1 是 NaN，所以最大尾数是"全 1 减 1"。"""
        m = self.man_bits
        frac = 2.0 - 2.0 ** (1 - m) if self.special == "fn" else 2.0 - 2.0 ** -m
        return 2.0 ** self.emax * frac

    @property
    def min_normal(self):
        """最小的正规正数 2^emin。"""
        return 2.0 ** self.emin

    @property
    def min_subnormal(self):
        """最小的正数（次正规）2^(emin − m)。比它的一半还小的数会被舍入成 0。"""
        return 2.0 ** (self.emin - self.man_bits)

    @property
    def eps(self):
        """机器精度（machine epsilon）：1 和"比 1 大的下一个数"之间的距离 = 2^−m。单位舍入误差 u = eps / 2。"""
        return 2.0 ** -self.man_bits

    def _default_saturate(self):
        return self.special == "none"

    # ---- 舍入 ----
    def round(self, x, rounding="nearest_even", saturate=None, rng=None):
        """
        把实数（numpy 数组或标量）舍入到本格式能表示的数，返回 float64 数组（值，而不是编码）。

        【思路】
            1. 找到 |x| 所在的"2 的幂区间" [2^e, 2^(e+1))：用 np.frexp 得到精确的 e（不用 log2，避免边界误差）。
               比 min_normal 小的数一律用 e = emin（次正规区的间隔和最小正规区一样）。
            2. 这个区间里相邻两个数的距离（ulp）是 q = 2^(e − m)。r = |x| / q 是一个"以 ulp 为单位"的实数，
               除以 2 的幂在 float64 里是精确的。
            3. 按舍入方式把 r 变成整数，再乘回 q。进位可能让结果跨进下一个区间（1.97 → 2.0），没关系，2.0 也能表示。
            4. 处理溢出：结果超过 max_value 时，饱和（截到 max_value）或变成 Inf / NaN。
        【参数】
            rounding："nearest_even"（就近舍入，平局取偶数——IEEE 默认，也是 numpy 的行为）、
                      "toward_zero"（截断）、"nearest_away"（就近，平局远离 0——小学的"四舍五入"）、
                      "stochastic"（随机舍入：以"离下面那个数的距离占 ulp 的比例"为概率向上舍入，期望无偏）。
            saturate：溢出时是否截到 max_value。None 表示用格式的默认：ieee / fn 不饱和（得到 Inf / NaN，
                      这是 IEEE 和 OCP 的 "OVF" 方式），none（FP6 / FP4）只能饱和。
            rng：随机舍入用的 numpy Generator（不传则用种子 0 新建一个）。
        【返回】float64 数组，形状和 x 相同；负数保留符号，负的小数舍入到 0 时得到 −0.0（和 IEEE 一致）。
        【边界】NaN 输入输出 NaN；±Inf 输入：ieee 不饱和时保持 Inf，其余情况截到 ±max_value（fn 不饱和时为 NaN）。
        【Python 小知识】np.frexp(a) 返回 (f, k) 使 a = f × 2^k、0.5 ≤ f < 1，所以 floor(log2 a) = k − 1，而且是精确的；
            np.ldexp(1.0, k) = 2^k；np.rint 是"就近、平局取偶"的取整。
        """
        if rounding not in ROUNDING_MODES:
            raise ValueError(f"rounding 只能是 {ROUNDING_MODES}")
        if saturate is None:
            saturate = self._default_saturate()
        x = np.asarray(x, dtype=np.float64)
        a = np.abs(x)
        finite = np.isfinite(a)
        a_f = np.where(finite, a, 0.0)
        _, k = np.frexp(a_f)
        e = np.maximum(k - 1, self.emin)
        q = np.ldexp(1.0, e - self.man_bits)
        r = a_f / q
        if rounding == "nearest_even":
            n = np.rint(r)
        elif rounding == "toward_zero":
            n = np.floor(r)
        elif rounding == "nearest_away":
            n = np.floor(r + 0.5)
        else:
            if rng is None:
                rng = np.random.default_rng(0)
            fl = np.floor(r)
            n = fl + (rng.random(r.shape) < (r - fl))
        y = n * q
        over = y > self.max_value
        if saturate or self.special == "none":
            y = np.where(over, self.max_value, y)
            y = np.where(np.isinf(a), self.max_value, y)
        else:
            overflow_value = np.inf if self.special == "ieee" else np.nan
            y = np.where(over | np.isinf(a), overflow_value, y)
        y = np.where(np.isnan(a), np.nan, y)
        return np.copysign(y, x)

    # ---- 编码 / 解码 ----
    def encode(self, x, rounding="nearest_even", saturate=None, rng=None):
        """
        把实数舍入并编码成比特，返回 int64 编码数组（0 … 2^bits − 1）。

        【思路】先 round() 得到能表示的值 v，再拆成三个字段：
            v = 0           → E = 0, M = 0
            v < min_normal  → 次正规：E = 0, M = v / min_subnormal（一定是整数）
            其它有限值       → E = e + bias, M = (v / 2^e − 1) × 2^m
            Inf             → E 全 1, M = 0（只有 ieee 格式）
            NaN             → ieee：E 全 1, M = 最高位为 1（"quiet NaN"）；fn：E 全 1, M 全 1
        【边界】"none" 格式没有 NaN 编码，输入里有 NaN 时抛出 ValueError。
        """
        v = self.round(x, rounding=rounding, saturate=saturate, rng=rng)
        m, top = self.man_bits, (1 << self.exp_bits) - 1
        sign = np.signbit(v).astype(np.int64)
        a = np.abs(v)
        E = np.zeros(a.shape, dtype=np.int64)
        M = np.zeros(a.shape, dtype=np.int64)
        nan = np.isnan(a)
        inf = np.isinf(a)
        if nan.any() and self.special == "none":
            raise ValueError(f"{self.name} 没有 NaN 编码")
        fin = ~(nan | inf)
        sub = fin & (a < self.min_normal)
        nor = fin & ~sub
        M[sub] = np.rint(a[sub] / self.min_subnormal).astype(np.int64)
        _, k = np.frexp(a[nor])
        e = k - 1
        E[nor] = e + self.bias
        M[nor] = np.rint((a[nor] / np.ldexp(1.0, e) - 1.0) * (1 << m)).astype(np.int64)
        E[inf] = top
        E[nan] = top
        M[nan] = (1 << m) - 1 if self.special == "fn" else 1 << (m - 1)
        return (sign << (self.exp_bits + m)) | (E << m) | M

    def decode(self, code):
        """
        把编码（整数或整数数组）还原成 float64 值。

        【思路】按 FloatFormat 的文档字符串里的公式逐字段计算；特殊值单独处理。
        【边界】编码超出 bits 位时只看低 bits 位。
        """
        code = np.asarray(code, dtype=np.int64) & ((1 << self.bits) - 1)
        m, top = self.man_bits, (1 << self.exp_bits) - 1
        S = (code >> (self.exp_bits + m)) & 1
        E = (code >> m) & top
        M = code & ((1 << m) - 1)
        val = np.where(E == 0,
                       M * self.min_subnormal,
                       np.ldexp(1.0 + M / float(1 << m), E - self.bias))
        if self.special == "ieee":
            val = np.where((E == top) & (M == 0), np.inf, val)
            val = np.where((E == top) & (M != 0), np.nan, val)
        elif self.special == "fn":
            val = np.where((E == top) & (M == (1 << m) - 1), np.nan, val)
        return np.where(S == 1, -val, val)

    def fields(self, code):
        """把一个编码拆成 (S, E, M) 三个整数（打印用）。"""
        code = int(code) & ((1 << self.bits) - 1)
        m = self.man_bits
        return code >> (self.exp_bits + m), (code >> m) & ((1 << self.exp_bits) - 1), code & ((1 << m) - 1)

    def all_values(self):
        """
        列出所有编码对应的值（长度 2^bits 的数组，下标就是编码）。

        【边界】只对 bits ≤ 16 开放（FP16 有 65,536 个编码；FP32 有 40 亿个，太多了）。
        """
        if self.bits > 16:
            raise ValueError("all_values 只用于 16 位及以下的格式")
        return self.decode(np.arange(1 << self.bits))

    def describe(self):
        """一行文字描述：E?M? 和 special 的含义。"""
        sp = {"ieee": "有 Inf/NaN", "fn": "无 Inf，1 个 NaN 编码", "none": "无 Inf/NaN"}[self.special]
        return f"{self.name}（E{self.exp_bits}M{self.man_bits}，bias {self.bias}，{sp}）"


# 预置格式 --------------------------------------------------------------------
FP32 = FloatFormat("FP32", 8, 23, 127, "ieee")
TF32 = FloatFormat("TF32", 8, 10, 127, "ieee")        # NVIDIA Tensor Core 的输入格式：FP32 的范围 + FP16 的精度（19 位，按 32 位存）
FP16 = FloatFormat("FP16", 5, 10, 15, "ieee")
BF16 = FloatFormat("BF16", 8, 7, 127, "ieee")         # "brain float"：FP32 的高 16 位
FP8_E4M3 = FloatFormat("FP8-E4M3", 4, 3, 7, "fn")     # OCP OFP8：权重、激活常用
FP8_E5M2 = FloatFormat("FP8-E5M2", 5, 2, 15, "ieee")  # OCP OFP8：梯度常用（范围大）
FP6_E3M2 = FloatFormat("FP6-E3M2", 3, 2, 3, "none")   # OCP MX
FP6_E2M3 = FloatFormat("FP6-E2M3", 2, 3, 1, "none")   # OCP MX
FP4_E2M1 = FloatFormat("FP4-E2M1", 2, 1, 1, "none")   # OCP MX；MXFP4 / NVFP4 的元素格式

FORMATS = [FP32, TF32, FP16, BF16, FP8_E4M3, FP8_E5M2, FP6_E3M2, FP6_E2M3, FP4_E2M1]
BY_NAME = {f.name: f for f in FORMATS}


def bf16_bit_trick(x32):
    """
    经典的 FP32 → BF16 就近舍入（平局取偶）的位运算写法，用来和 BF16.round 对照。

    【思路】BF16 就是 FP32 的高 16 位。要就近舍入，先给 FP32 的位模式加上 0x7FFF（"差一点到一半"），
        再加上"保留部分的最低位"（平局时向偶数靠），然后右移 16 位丢掉低位。
    【参数】x32：np.float32 数组。【返回】float32 数组（高 16 位是 BF16，低 16 位为 0）。
    【边界】不处理 NaN（NaN 加上偏置可能变成 Inf），调用方要先去掉 NaN。
    【Python 小知识】arr.view(np.uint32) 不改变内存，只是"换一种类型来看"同一串比特。
    """
    u = np.asarray(x32, dtype=np.float32).view(np.uint32).astype(np.uint64)
    u = (u + 0x7FFF + ((u >> 16) & 1)) >> 16
    return (u << 16).astype(np.uint32).view(np.float32)


# ---------------------------------------------------------------------------
# 3. 块缩放：MX 与 NVFP4
# ---------------------------------------------------------------------------
def mx_quantize(x, elem, block=32):
    """
    按 OCP MX 规范 §6.3 把一维数组量化成 MX 格式，再反量化，返回 (近似值, 每块的缩放因子)。

    【思路】每 block 个数一块，共享一个 E8M0 缩放因子 X（只能是 2 的幂）：
        X = 2^floor(log2(块内最大绝对值)) / 2^(元素格式的 emax)
        （"块内最大值所在的 2 的幂"除以"元素格式能表示的最大的 2 的幂"），
        元素 P = round(V / X)，超出元素格式最大值的截到最大值。最后近似值 = P × X。
    【参数】x：一维数组（长度必须是 block 的倍数）；elem：元素格式（FloatFormat）；block：块大小（规范是 32）。
    【边界】全零的块：缩放因子取 1，元素全为 0。E8M0 的指数范围是 [−127, 127]，超出时截断。
        这种"按 2 的幂向下取"的缩放让块内最大值除完以后落在 [2^emax, 2^(emax+1))，可能超过元素的最大值
        （例如 FP4 最大 6，而 7.9 也在 [4, 8) 里），这些数会被截到最大值——这是 MX 规范的已知代价。
    """
    x = np.asarray(x, dtype=np.float64)
    if x.size % block:
        raise ValueError("长度必须是 block 的倍数")
    xb = x.reshape(-1, block)
    amax = np.abs(xb).max(axis=1, keepdims=True)
    _, k = np.frexp(np.where(amax > 0, amax, 1.0))
    shared = np.clip(k - 1 - elem.emax, -127, 127)
    X = np.ldexp(1.0, shared)
    P = elem.round(xb / X, saturate=True)
    return (P * X).reshape(x.shape), X.ravel()


def nvfp4_quantize(x, block=16):
    """
    NVFP4：每 16 个数一块，块缩放因子用 FP8 E4M3 存；整个张量再有一个 FP32 缩放因子 g。返回 (近似值, 块缩放, g)。

    【思路】
        g = max|x| / (6 × 448)：把"块缩放因子"挪到 E4M3 的表示范围里（块内最大值 / 6 最大是 max|x| / 6，除以 g 后最大 448）；
        块缩放 s_b = round_E4M3(块内最大值 / 6 / g)：E4M3 有 3 位尾数，不必是 2 的幂，能让块内最大值几乎正好落在 6；
        元素 = round_E2M1(V / (s_b × g))，近似值 = 元素 × s_b × g。
    【边界】长度必须是 block 的倍数；全零的张量或块直接输出 0。s_b 是舍入过的，块内最大值除完可能略大于 6，截到 6。
    """
    x = np.asarray(x, dtype=np.float64)
    if x.size % block:
        raise ValueError("长度必须是 block 的倍数")
    amax = np.abs(x).max()
    if amax == 0:
        return np.zeros_like(x), np.zeros(x.size // block), 0.0
    g = float(FP32.round(amax / (FP4_E2M1.max_value * FP8_E4M3.max_value)))
    xb = x.reshape(-1, block)
    bmax = np.abs(xb).max(axis=1, keepdims=True)
    sb = FP8_E4M3.round(bmax / FP4_E2M1.max_value / g, saturate=True)
    scale = sb * g
    safe = np.where(scale > 0, scale, 1.0)
    q = np.where(scale > 0, FP4_E2M1.round(xb / safe) * scale, 0.0)
    return q.reshape(x.shape), sb.ravel(), g


def tensor_scaled(x, elem):
    """整个张量一个缩放因子（FP32）：s = max|x| / elem.max_value，近似值 = round(x / s) × s。全零时原样返回。"""
    x = np.asarray(x, dtype=np.float64)
    amax = np.abs(x).max()
    if amax == 0:
        return x.copy()
    s = amax / elem.max_value
    return elem.round(x / s, saturate=True) * s


# ---------------------------------------------------------------------------
# 4. 整数量化（仿射量化）
# ---------------------------------------------------------------------------
def quantize_int(x, bits=8, symmetric=True, granularity="tensor", group=None, clip=None):
    """
    整数量化再反量化：x ≈ scale × (q − zero)，q 是 bits 位整数。返回 (近似值, q, scale, zero)。

    【两种映射】
        对称（symmetric）：zero = 0，q ∈ [−(2^(b−1) − 1), 2^(b−1) − 1]（不用 −128，正负对称），
            scale = max|x| / (2^(b−1) − 1)。适合以 0 为中心的数据（权重）。
        非对称（asymmetric）：q ∈ [0, 2^b − 1]，scale = (max − min) / (2^b − 1)，zero = round(−min / scale)。
            适合偏向一侧的数据（ReLU 之后全是非负数）：所有 2^b 个编码都用上了。
    【粒度 granularity】
        "tensor"：整个数组一个 scale；
        "channel"：x 是二维 (行, 列)，每一行一个 scale（权重矩阵按输出通道）；
        "group"：x 的最后一维每 group 个数一个 scale（例如 group = 128，GPTQ / AWQ 的常用设置）。
    【clip】给定时，把范围缩小到 clip × max|x|（对称）或 clip × (max − min)（非对称，按比例缩两端），
        用"截掉少数大值"换"其余的数更精细"（05 脚本演示 2）。
    【边界】范围为 0（整组都是同一个数）时 scale 取 1，避免除以 0；group 必须能整除最后一维。
    """
    x = np.asarray(x, dtype=np.float64)
    if granularity == "tensor":
        xs = x.reshape(1, -1)
    elif granularity == "channel":
        if x.ndim != 2:
            raise ValueError("channel 粒度需要二维数组（行 = 通道）")
        xs = x
    elif granularity == "group":
        if group is None or x.shape[-1] % group:
            raise ValueError("group 必须给出且能整除最后一维")
        xs = x.reshape(-1, group)
    else:
        raise ValueError("granularity 只能是 tensor / channel / group")
    c = 1.0 if clip is None else clip
    if symmetric:
        qmax = (1 << (bits - 1)) - 1
        amax = np.abs(xs).max(axis=1, keepdims=True) * c
        scale = np.where(amax > 0, amax / qmax, 1.0)
        zero = np.zeros_like(scale)
        q = np.clip(np.rint(xs / scale), -qmax, qmax)
    else:
        qmax = (1 << bits) - 1
        lo = xs.min(axis=1, keepdims=True)
        hi = xs.max(axis=1, keepdims=True)
        mid = (lo + hi) / 2
        lo, hi = mid + (lo - mid) * c, mid + (hi - mid) * c
        scale = np.where(hi > lo, (hi - lo) / qmax, 1.0)
        zero = np.clip(np.rint(-lo / scale), 0, qmax)
        q = np.clip(np.rint(xs / scale) + zero, 0, qmax)
    deq = (scale * (q - zero)).reshape(x.shape)
    return deq, q.reshape(x.shape).astype(np.int64), scale, zero


# ---------------------------------------------------------------------------
# 5. 低精度累加
# ---------------------------------------------------------------------------
def accumulate(values, acc=None, method="sequential", chunk=None, rounding="nearest_even", seed=0):
    """
    把二维数组 values（每一行是一组要相加的数）逐行求和，模拟"累加器只有 acc 这种精度"的硬件。

    【参数】
        acc：累加器格式（FloatFormat）；None 表示用 float64（作为"几乎精确"的参照）。
        method：
            "sequential"：s ← round(s + x_j)，一个一个加（最常见，误差最大）；
            "pairwise"：两两相加、再两两相加……（树形归约，GPU 的并行归约天然就是这样，误差小得多）；
            "kahan"：Kahan 补偿求和——用另一个变量 c 记住每次加法丢掉的低位，下一次补回来。
        chunk：给定时，每 chunk 个数先在 acc 里累加，再把这段的部分和加到一个 FP32 累加器里
            （DeepSeek-V3 的"每 N_C = 128 个元素提升到 CUDA core 的 FP32 寄存器"）。只和 sequential 一起用。
        rounding：每次加法后的舍入方式（见 FloatFormat.round）。
        seed：随机舍入的种子。
    【返回】一维 float64 数组，长度 = 行数。
    【思路】所有行同时算（numpy 向量化），列方向用 Python 循环——所以 n 列就有 n 次循环，n 到几万也只要零点几秒。
        两个能用 acc 精确表示的数相加，先在 float64 里算、再舍入到 acc，结果和"真正的 acc 硬件加法"相同
        （float64 的 53 位 ≥ 2p + 2，双重舍入不会出错，p ≤ 24）。
    【边界】values 里的数最好已经能被 acc 表示（否则第一次加法前就有一次舍入）；method 不认识时抛出 ValueError。
    """
    v = np.asarray(values, dtype=np.float64)
    if v.ndim == 1:
        v = v[None, :]
    rng = np.random.default_rng(seed)

    def rnd(z):
        return z if acc is None else acc.round(z, rounding=rounding, rng=rng)

    T, n = v.shape
    if method == "sequential":
        if chunk is None:
            s = np.zeros(T)
            for j in range(n):
                s = rnd(s + v[:, j])
            return s
        outer = np.zeros(T)
        for start in range(0, n, chunk):
            s = np.zeros(T)
            for j in range(start, min(start + chunk, n)):
                s = rnd(s + v[:, j])
            outer = FP32.round(outer + s)
        return outer
    if method == "pairwise":
        cur = v.copy()
        while cur.shape[1] > 1:
            if cur.shape[1] % 2:
                cur = np.concatenate([cur, np.zeros((T, 1))], axis=1)
            cur = rnd(cur[:, 0::2] + cur[:, 1::2])
        return cur[:, 0]
    if method == "kahan":
        s = np.zeros(T)
        c = np.zeros(T)
        for j in range(n):
            y = rnd(v[:, j] - c)
            t = rnd(s + y)
            c = rnd(rnd(t - s) - y)
            s = t
        return s
    raise ValueError("method 只能是 sequential / pairwise / kahan")


# ---------------------------------------------------------------------------
# 6. softmax 与方差：数值稳定的写法
# ---------------------------------------------------------------------------
def softmax_naive(x, dtype=np.float32):
    """直接按定义 exp(x) / Σexp(x) 计算（在 dtype 里）。x 很大时 exp 溢出成 Inf，Inf / Inf = NaN。"""
    x = np.asarray(x, dtype=dtype)
    with np.errstate(over="ignore", invalid="ignore"):
        e = np.exp(x)
        return e / e.sum(dtype=dtype)


def softmax_safe(x, dtype=np.float32):
    """先减去最大值再 exp：数学上结果不变（分子分母同乘 e^(−max)），但 exp 的参数 ≤ 0，不会溢出。"""
    x = np.asarray(x, dtype=dtype)
    e = np.exp(x - x.max())
    return e / e.sum(dtype=dtype)


def online_softmax_stats(x, dtype=np.float32):
    """
    在线（online）softmax 的"一遍扫描"：同时维护目前见过的最大值 m 和以 m 为基准的指数和 d。

    【思路】（Milakov & Gimelshein 2018）每来一个 x_j：
        m_new = max(m, x_j)
        d     = d × exp(m − m_new) + exp(x_j − m_new)      ← 旧的和要"换一个基准"
        扫完以后 softmax_j = exp(x_j − m) / d。
        这样求最大值和求和合并成一遍，少读一次输入（论文：每个元素的访存从 4 次降到 3 次）。
    【返回】(m, d)，都是 dtype 类型的标量。
    【边界】x 不能是空数组；全是 −Inf 时 exp(−Inf − (−Inf)) = NaN（真实内核会特判，这里不处理）。
    """
    x = np.asarray(x, dtype=dtype)
    m = dtype(-np.inf)
    d = dtype(0.0)
    for xj in x:
        m_new = max(m, xj)
        d = d * np.exp(m - m_new) + np.exp(xj - m_new)
        m = m_new
    return dtype(m), dtype(d)


def merge_softmax_stats(a, b):
    """
    合并两段的 (m, d)：m = max(m_a, m_b)，d = d_a × exp(m_a − m) + d_b × exp(m_b − m)。

    【系统联系】这个"⊕"运算满足结合律，所以一行可以切成很多段、由不同的线程块各算一段，最后再合并——
        这就是 FlashAttention 分块计算 softmax、flash-decoding 把 KV 切成几段并行的数学基础（第 11 讲）。
    """
    (ma, da), (mb, db) = a, b
    m = max(ma, mb)
    return m, da * np.exp(ma - m) + db * np.exp(mb - m)


def variance_one_pass(x, dtype=np.float32):
    """一遍公式 Var = E[x²] − (E[x])²（在 dtype 里累加）。均值远大于标准差时会发生灾难性抵消。"""
    x = np.asarray(x, dtype=dtype)
    n = dtype(x.size)
    s = x.sum(dtype=dtype)
    s2 = (x * x).sum(dtype=dtype)
    return s2 / n - (s / n) * (s / n)


def variance_two_pass(x, dtype=np.float32):
    """两遍公式：先求均值，再求 Σ(x − mean)² / n。要读两遍数据，但数值上稳定得多。"""
    x = np.asarray(x, dtype=dtype)
    mu = x.sum(dtype=dtype) / dtype(x.size)
    d = x - mu
    return (d * d).sum(dtype=dtype) / dtype(x.size)


def variance_welford(x, dtype=np.float32):
    """
    Welford 算法：一遍扫描、数值稳定。维护 n、均值 mean 和"离差平方和" M2：
        delta = x − mean；mean += delta / n；M2 += delta × (x − mean)。最后 Var = M2 / n。
    【系统联系】LayerNorm / RMSNorm 的融合内核想"只读一遍激活"，用的就是这类一遍且稳定的算法
        （以及它的并行合并版本，和 merge_softmax_stats 是同一个思路）。
    【边界】逐元素 Python 循环，只适合几万个数以内的演示。
    """
    x = np.asarray(x, dtype=dtype)
    mean = dtype(0.0)
    m2 = dtype(0.0)
    for i, xi in enumerate(x, start=1):
        delta = xi - mean
        mean = mean + delta / dtype(i)
        m2 = m2 + delta * (xi - mean)
    return m2 / dtype(x.size)


# ---------------------------------------------------------------------------
# 7. 误差指标
# ---------------------------------------------------------------------------
def rel_error(approx, exact):
    """相对误差 ||approx − exact|| / ||exact||（2-范数）。exact 全为 0 时返回 inf。"""
    approx = np.asarray(approx, dtype=np.float64)
    exact = np.asarray(exact, dtype=np.float64)
    den = np.linalg.norm(exact)
    return np.inf if den == 0 else float(np.linalg.norm(approx - exact) / den)


def sqnr_db(approx, exact):
    """信号量化噪声比 SQNR = 20·log10(||exact|| / ||approx − exact||)，单位 dB，越大越好。每多 1 位大约 +6 dB。"""
    r = rel_error(approx, exact)
    return np.inf if r == 0 else float(-20 * np.log10(r))


# ---------------------------------------------------------------------------
# 自检
# ---------------------------------------------------------------------------
def _selftest():
    print("== numfmt 自检 ==")
    # 1) 规格表里的关键数字（OCP / IEEE 的表格）
    expect = {"FP32": 3.4028234663852886e38, "FP16": 65504.0, "BF16": 3.3895313892515355e38,
              "FP8-E4M3": 448.0, "FP8-E5M2": 57344.0, "FP6-E3M2": 28.0, "FP6-E2M3": 7.5, "FP4-E2M1": 6.0}
    for name, mx in expect.items():
        assert BY_NAME[name].max_value == mx, (name, BY_NAME[name].max_value)
    assert FP8_E4M3.min_subnormal == 2.0 ** -9 and FP8_E5M2.min_subnormal == 2.0 ** -16
    assert FP4_E2M1.min_subnormal == 0.5 and FP6_E3M2.min_subnormal == 0.0625 and FP6_E2M3.min_subnormal == 0.125
    vals = FP4_E2M1.all_values()
    assert sorted(set(np.abs(vals).tolist())) == [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
    print("  规格表（最大值、最小次正规数、FP4 的 8 个非负值）：通过")

    # 2) 和 numpy 的硬件实现逐位比较
    rng = np.random.default_rng(1)
    x = np.concatenate([rng.standard_normal(50_000) * np.exp2(rng.integers(-30, 20, 50_000)),
                        rng.standard_normal(1000) * 1e-6, [0.0, -0.0, 65504.0, 65519.9, 65520.0, 1e6, -1e6]])
    ours = FP16.encode(x)
    with np.errstate(over="ignore"):
        ref = x.astype(np.float16).view(np.uint16).astype(np.int64)
    assert np.array_equal(ours, ref), "FP16 编码和 numpy 不一致"
    x32 = rng.standard_normal(50_000) * np.exp2(rng.integers(-140, 120, 50_000).astype(float))
    with np.errstate(over="ignore"):
        ref32 = x32.astype(np.float32).view(np.uint32).astype(np.int64)
    assert np.array_equal(FP32.encode(x32), ref32), "FP32 编码和 numpy 不一致"
    print("  FP16 / FP32 编码和 numpy 逐位一致（各 5 万个数，含次正规数和溢出）：通过")
    f32 = rng.standard_normal(50_000).astype(np.float32) * np.float32(1000)
    assert np.array_equal(BF16.round(f32), bf16_bit_trick(f32).astype(np.float64)), "BF16 和位运算写法不一致"
    print("  BF16 和'加 0x7FFF 再右移'的位运算写法一致：通过")

    # 3) 编码 → 解码 往返
    for f in (FP8_E4M3, FP8_E5M2, FP6_E3M2, FP6_E2M3, FP4_E2M1, FP16):
        codes = np.arange(1 << f.bits)
        v = f.decode(codes)
        ok = np.isfinite(v)
        back = f.encode(v[ok])
        assert np.array_equal(back, codes[ok]), f.name
    print("  所有有限编码 解码→编码 往返一致：通过")

    # 4) 块缩放和整数量化的基本性质
    y = rng.standard_normal(4096)
    q, X = mx_quantize(y, FP4_E2M1)
    assert np.all(np.log2(X) == np.round(np.log2(X))), "MX 缩放因子必须是 2 的幂"
    qn, sb, g = nvfp4_quantize(y)
    assert rel_error(qn, y) < rel_error(q, y) < 0.2
    d, qi, s, z = quantize_int(y, 8)
    assert qi.min() >= -127 and qi.max() <= 127 and rel_error(d, y) < 0.01
    print(f"  MXFP4 相对误差 {rel_error(q, y):.3f} > NVFP4 {rel_error(qn, y):.3f}；INT8 对称 {rel_error(d, y):.4f}：通过")

    # 5) 累加：BF16 里一直加 1，停在 256
    ones = np.ones((1, 300))
    assert accumulate(ones, BF16)[0] == 256.0 and accumulate(ones, BF16, method="pairwise")[0] == 300.0
    print("  BF16 顺序累加 300 个 1 = 256，两两相加 = 300：通过")
    m, dd = online_softmax_stats(np.array([1.0, 3.0, 2.0]))
    assert abs(np.exp(3.0 - m) / dd - softmax_safe([1.0, 3.0, 2.0])[1]) < 1e-6
    print("  在线 softmax 与两遍 softmax 一致：通过")


if __name__ == "__main__":
    _selftest()