"""
=============================================================================
CS349E · 前置讲 P3 · 01_integers_and_fixed_point.py —— 整数、补码、溢出，以及为什么需要浮点数
=============================================================================

【这个脚本要回答的问题】
    计算机里的数都是有限的比特。最简单的是整数：8 位整数只有 256 个编码。本脚本回答：

      演示 1：负数怎样用比特表示？为什么几乎所有硬件都用"补码"？
      演示 2：溢出了会怎样？回绕（127 + 1 = −128）和饱和（127 + 1 = 127）各在什么场合用？
      演示 3：INT8 量化的矩阵乘法，累加器为什么必须是 32 位？（真的用 numpy 算 1000 个长度 4096 的点积）
      演示 4：定点数（固定小数点）能表示多大范围？为什么深度学习需要"浮点"——让小数点随数值大小移动？

【怎么读输出】全部是确定性的计算（固定随机种子），每次运行结果相同。

【运行】python3 01_integers_and_fixed_point.py      （不到 1 秒）
【依赖】numpy；从 ../lecture01 导入 console（表格打印）；同目录的 numfmt。
"""

import os
import sys

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

from console import print_header, print_observations, print_table  # noqa: E402
from numfmt import (BF16, FP16, bits_str, int_range, saturate_int,  # noqa: E402
                    twos_complement_decode, twos_complement_encode, wrap_int)


# ---------------------------------------------------------------------------
# 演示 1：补码
# ---------------------------------------------------------------------------
def demo1_twos_complement():
    """
    列出几个 8 位整数的补码编码，并验证"取反加一"和"同一个加法器"两个性质。

    【补码的定义】编码 = v mod 2^8。所以 −1 的编码是 255（全 1），−128 的编码是 128（只有最高位是 1）。
        换一种看法：最高位的"权重"是 −128，其余位是 +64, +32, …, +1。
    【为什么用补码】
        1. 加法器不用区分有符号还是无符号：把两个编码当作无符号数相加、丢掉第 9 位，结果就是正确的补码。
        2. 0 只有一个编码（原码和反码都有 +0 和 −0 两个）。
        3. 取负数 = 按位取反再加 1，硬件很便宜。
    """
    print_header("演示 1：8 位补码——最高位的权重是 −128")
    rows = []
    for v in (0, 1, 2, 127, -1, -2, -127, -128):
        code = twos_complement_encode(v, 8)
        if v == -128:
            neg_check = "—（+128 放不下）"
        else:
            inv_plus_one = (~twos_complement_encode(-v, 8) + 1) & 0xFF
            neg_check = "✓" if inv_plus_one == code else "✗"
        rows.append([v, bits_str(code, 8, (1, 7)), code, twos_complement_decode(code, 8), neg_check])
    print_table(["整数 v", "补码（符号位 | 其余 7 位）", "当作无符号数", "解码回来", "v = ~(−v) + 1 ？"], rows)

    print("\n用同一个'无符号加法器'做有符号加法（结果只保留低 8 位）：")
    rows = []
    for a, b in ((-3, 5), (100, -100), (-50, -60), (100, 27)):
        ca, cb = twos_complement_encode(a, 8), twos_complement_encode(b, 8)
        s = (ca + cb) & 0xFF
        rows.append([f"{a} + {b}", f"{ca} + {cb} = {ca + cb}", f"{(ca + cb) & 0xFF}（mod 256）",
                     twos_complement_decode(s, 8), a + b])
    print_table(["有符号运算", "编码相加（无符号）", "丢掉进位", "解码", "数学上的答案"], rows)
    lo, hi = int_range(8)
    print_observations([
        f"8 位补码的范围是 [{lo}, {hi}]：负数比正数多一个，因为 0 占用了一个'非负'的编码。",
        "−128 是个怪数：它的相反数 +128 表示不了（取反加一后还是它自己）。所以对称量化通常不用 −128，"
        "只用 [−127, 127]——OCP 的 MX 规范也明确允许把它空着'以保持正负对称、避免负偏差'。",
        "四个加法都用无符号加法器算对了：硬件里有符号和无符号整数共用同一套加法器，这是补码胜出的主要原因。",
    ])


# ---------------------------------------------------------------------------
# 演示 2：溢出——回绕 vs 饱和
# ---------------------------------------------------------------------------
def demo2_overflow():
    """
    比较两种溢出处理：
        回绕（wraparound）：按 2^8 取模，C 语言和 numpy 的整数类型都是这样（不报错！）；
        饱和（saturation）：截到最大 / 最小值，DSP 和量化里常用（int8 量化时超范围的数截到 ±127）。
    【Python 小知识】numpy 的 int8 数组做运算溢出时**静默回绕**，不报错也不警告。
        np.abs 对 int8 的 −128 返回 −128（因为 +128 放不下）——这是一个真实存在的坑。
    """
    print_header("演示 2：溢出——回绕（wrap）还是饱和（saturate）？")
    rows = []
    for expr, val in (("127 + 1", 127 + 1), ("−128 − 1", -128 - 1), ("100 + 100", 200),
                      ("−100 × 2", -200), ("64 × 4", 256)):
        rows.append([expr, val, int(wrap_int(val, 8)), int(saturate_int(val, 8))])
    print_table(["运算（8 位）", "数学结果", "回绕的结果", "饱和的结果"], rows)

    a = np.array([127, -128, 100], dtype=np.int8)
    b = np.array([1, -1, 100], dtype=np.int8)
    print(f"\nnumpy 实测：int8 数组 {a.tolist()} + {b.tolist()} = {(a + b).tolist()}（静默回绕）")
    print(f"numpy 实测：np.abs(int8 数组 [-128, -5]) = {np.abs(np.array([-128, -5], dtype=np.int8)).tolist()}")
    print_observations([
        "回绕的结果离真实答案'很远'（127 + 1 变成 −128，符号都反了）；饱和的结果'最接近'能表示的值。",
        "所以量化（把实数变成 int8）时用饱和：一个超范围的权重被截到 127，误差有限；如果回绕成 −128，就是灾难。",
        "numpy 的 int8 / int16 溢出不会报错——写低精度代码时要自己检查范围。",
    ])


# ---------------------------------------------------------------------------
# 演示 3：INT8 点积的累加器宽度
# ---------------------------------------------------------------------------
def demo3_accumulator_width(T=1000, K=4096, seed=0):
    """
    INT8 矩阵乘法：输入是 int8，乘积最大 127 × 127 = 16,129，K 个乘积相加。累加器要多宽？

    【实验】随机生成 T 对长度 K 的 int8 向量（两种分布：均匀分布在 [−127, 127]；
        "量化过的高斯"——标准差约 32 的整数，更像真实的量化权重和激活），比较三种累加器：
        int16 回绕、int16 饱和、int32。精确答案用 int64 算。
    【一个微妙之处】回绕是模 2^16 的运算：即使中间的部分和溢出了，只要**最终结果**落在 int16 范围里，
        回绕的累加器给出的结果仍然是对的（模运算的性质）。饱和的累加器一旦中途截断，就再也回不来了。
    【Python 小知识】np.cumsum(p, axis=1, dtype=np.int16) 沿每一行做前缀和，并且每一步都在 int16 里（回绕）。
    """
    print_header(f"演示 3：INT8 点积（K = {K}）的累加器要多宽？——{T} 个点积实测")
    worst = 127 * 127
    lo16, hi16 = int_range(16)
    lo32, hi32 = int_range(32)
    print(f"理论：一个乘积最大 {worst:,}；int16 最多放下 {hi16 // worst} 个最坏的乘积，"
          f"int32 最多 {hi32 // worst:,} 个。")
    rng = np.random.default_rng(seed)
    gens = {
        "均匀分布 [−127, 127]": lambda: rng.integers(-127, 128, (T, K)),
        "量化过的高斯（σ ≈ 32）": lambda: np.clip(np.rint(rng.standard_normal((T, K)) * 127 / 4), -127, 127).astype(np.int64),
    }
    rows = []
    for name, gen in gens.items():
        a, b = gen(), gen()
        p = a * b
        exact = p.sum(axis=1)
        partial = np.cumsum(p, axis=1)
        wrap16 = np.cumsum(p.astype(np.int16), axis=1, dtype=np.int16)[:, -1].astype(np.int64)
        sat16 = np.zeros(T, dtype=np.int64)
        for j in range(K):
            sat16 = np.clip(sat16 + p[:, j], lo16, hi16)
        acc32 = np.cumsum(p.astype(np.int32), axis=1, dtype=np.int32)[:, -1].astype(np.int64)
        rows.append([name, f"{np.abs(exact).mean():,.0f}", f"{np.abs(partial).max():,}",
                     f"{np.mean(wrap16 != exact) * 100:.1f}%", f"{np.mean(sat16 != exact) * 100:.1f}%",
                     f"{np.mean(acc32 != exact) * 100:.1f}%"])
    print_table(["输入分布", "|结果| 的平均", "部分和的最大 |值|", "int16 回绕：错", "int16 饱和：错", "int32：错"], rows)
    print_observations([
        "长度 4096 的 INT8 点积，结果动辄几万到几十万，int16 根本放不下：两种分布下 int16 累加器都有一半以上的结果是错的。",
        "int32 一个都没错。这就是为什么所有 INT8 矩阵乘法硬件（GPU Tensor Core 的 s8 × s8 → s32、TPU 的 MXU）"
        "都把乘积累加到 32 位：输入省了，累加器不能省。",
        "回绕的 int16 比饱和的错得少——不是因为它更'聪明'，而是模运算让'中途溢出但最终落回范围'的结果碰巧对了。"
        "这个性质在硬件设计里真有人用（只要能证明最终结果不溢出），但不能指望。",
        "浮点累加有类似的问题，只是表现为'精度不够'而不是'溢出'（03 脚本演示 3：FP8 矩阵乘法的累加精度）。",
    ])


# ---------------------------------------------------------------------------
# 演示 4：定点数 vs 浮点数
# ---------------------------------------------------------------------------
def demo4_fixed_vs_float(n=10_000, seed=0):
    """
    同样 16 位，定点数（Q7.8：1 位符号、7 位整数、8 位小数，间隔 2^−8 ≈ 0.0039、最大约 128）
    和浮点数（FP16、BF16）在不同数量级上的**最大相对误差**。

    【思路】对每个十进制数量级 [10^k, 10^(k+1))，取 n 个对数均匀分布的数，分别舍入到三种格式，算最大相对误差。
    【为什么看相对误差】神经网络关心的是"这个数大约是多少"，0.001 错了 0.0039 是 400%，1000 错了 0.0039 无所谓。
    【边界】定点数超出范围时饱和到 ±127.996；FP16 超过 65,504 变成 Inf（相对误差记为 inf）。
    """
    print_header("演示 4：同样 16 位——定点数 Q7.8 vs FP16 vs BF16，各数量级上的最大相对误差")
    rng = np.random.default_rng(seed)
    rows = []
    for k in range(-5, 5):
        v = np.exp(rng.uniform(np.log(10.0 ** k), np.log(10.0 ** (k + 1)), n))
        fx = np.clip(np.rint(v * 256), -32768, 32767) / 256
        r_fx = np.max(np.abs(fx - v) / v)
        with np.errstate(invalid="ignore"):
            r16 = np.max(np.abs(FP16.round(v) - v) / v)
        rb = np.max(np.abs(BF16.round(v) - v) / v)
        rows.append([f"[1e{k}, 1e{k + 1})", f"{r_fx:.1e}", f"{r16:.1e}", f"{rb:.1e}"])
    print_table(["数量级", "定点 Q7.8", "FP16", "BF16"], rows)
    print_observations([
        "定点数的**绝对**误差是常数（间隔 2^−8），所以只在 1–100 附近好用：小于 0.01 的数相对误差接近 100%（几乎都变成 0），"
        "大于 128 的数被截断。",
        f"FP16 在 1e−4 到 6 万之间的相对误差都约为 4.9e−4 = 2^−11（半个 ulp）：浮点数让'小数点'跟着数值大小移动，"
        "把相对精度均匀地分配到很宽的范围上。小于约 6e−5 时进入次正规区，精度开始下降；超过 65,504 溢出成 Inf。",
        "BF16 的相对误差大 8 倍（3.9e−3 = 2^−8），但在整个表里都没有溢出——它的指数和 FP32 一样有 8 位。"
        "'范围'和'精度'是两种不同的资源，16 个比特怎么分给它们，是 02 脚本的主题。",
        "LLM 的权重、激活、梯度跨越很多个数量级，所以训练和推理都用浮点数；只有在'先缩放到一个已知范围'之后"
        "（量化，05 脚本），整数才重新登场。",
    ])


if __name__ == "__main__":
    demo1_twos_complement()
    demo2_overflow()
    demo3_accumulator_width()
    demo4_fixed_vs_float()