"""
mem_calc.py —— 模型的"账本"：参数量、显存、FLOPs、时间与花费
=====================================================================

【这个文件要回答的问题】
    在租 GPU 之前，你应该能用纸笔（或者这个文件）回答：
      1. 这个模型有多少参数？——由结构超参数精确算出（和 PyTorch 逐个核对过）。
      2. 推理要多少显存？——权重 + KV 缓存。
      3. 训练要多少显存？——权重 + 梯度 + 优化器状态 + 激活值。
      4. 训练要算多少次浮点运算？——约 6 × 参数量 × token 数（"6ND"）。
      5. 要多少 GPU 小时、多少钱？——FLOPs ÷ (GPU 峰值 × 利用率)，再乘单价。
    这些估算决定了：用什么卡、要几张、要不要换方法（LoRA、量化、ZeRO……），以及这个实验做不做得起。

【怎样使用】
    from mem_calc import PRESETS, count_params, training_memory, kv_cache_bytes, fmt_bytes
    cfg = PRESETS["llama3-8b"]
    count_params(cfg)                      # 8,030,261,248
    fmt_bytes(training_memory(cfg)["total"])

【边界】
    - 激活值的显存和实现细节（是否用 FlashAttention、是否有 dropout、是否重算）关系很大，
      这里只给量级估计（Korthikanti et al., 2022 的公式），精确值要实测（见 02 号脚本）。
    - GPU 的峰值算力取的是公开规格里的稠密 BF16 张量核心算力；实际能达到的比例（MFU）
      通常在 30%～60%，是一个需要你自己假设或实测的数。
    - 价格随云厂商、地区、时间变化很大，本文件**不内置任何价格**，一律由调用者传入。
"""

from dataclasses import dataclass, replace

DTYPE_BYTES = {"fp32": 4, "fp16": 2, "bf16": 2, "fp8": 1, "int8": 1, "int4": 0.5}

GIB = 1024 ** 3


@dataclass(frozen=True)
class Config:
    """
    decoder-only Transformer 的结构超参数。

    【字段】vocab：词表大小 V；d_model：隐藏维度 d；n_layers：层数 L；n_heads：注意力头数；
            n_kv_heads：K/V 头数（= n_heads 是普通多头注意力，更少就是 GQA）；d_ff：MLP 中间层宽度；
            mlp："gelu"（2 个矩阵）或 "swiglu"（3 个矩阵）；norm："layernorm" 或 "rmsnorm"；
            bias：线性层是否有偏置；tied：输出层是否与输入嵌入共享；max_ctx：最大上下文长度；
            learned_pos：是否有可学习的位置嵌入（max_ctx × d 个参数）。
    【Python 小知识】@dataclass 自动生成 __init__ 等方法；frozen=True 让对象不可修改（可以当字典的键）。
            dataclasses.replace(cfg, n_layers=4) 得到一个只改了某个字段的新对象。
    """
    name: str
    vocab: int
    d_model: int
    n_layers: int
    n_heads: int
    n_kv_heads: int
    d_ff: int
    mlp: str = "swiglu"
    norm: str = "rmsnorm"
    bias: bool = False
    tied: bool = False
    max_ctx: int = 4096
    learned_pos: bool = False

    @property
    def head_dim(self):
        return self.d_model // self.n_heads


# 几个公开模型的结构（来自它们公开的配置文件）
PRESETS = {
    "gpt2-small": Config("GPT-2 small", 50257, 768, 12, 12, 12, 3072, mlp="gelu", norm="layernorm",
                         bias=True, tied=True, max_ctx=1024, learned_pos=True),
    "llama2-7b": Config("Llama 2 7B", 32000, 4096, 32, 32, 32, 11008, max_ctx=4096),
    "llama3-8b": Config("Llama 3 8B", 128256, 4096, 32, 32, 8, 14336, max_ctx=8192),
    "llama2-70b": Config("Llama 2 70B", 32000, 8192, 80, 64, 8, 28672, max_ctx=4096),
}


def param_breakdown(c):
    """
    按部件拆分的参数量。

    【公式】（d = d_model，hd = head_dim）
        注意力：Q 和 O 各 d × (头数·hd)；K、V 各 d × (KV 头数·hd)；有偏置时再加上输出维度。
        MLP：gelu 是 2 个 d × d_ff 矩阵，swiglu 是 3 个；有偏置时加 d_ff（每个上投影）和 d（下投影）。
        归一化：layernorm 每处 2d（缩放 + 偏置），rmsnorm 每处 d；每层 2 处，最后还有 1 处。
        嵌入：V × d；位置嵌入：max_ctx × d；输出层：不共享时再加 V × d。
    【返回】dict，各部件的参数个数和 total。
    """
    d, hd = c.d_model, c.head_dim
    attn = d * c.n_heads * hd * 2 + d * c.n_kv_heads * hd * 2
    if c.bias:
        attn += c.n_heads * hd + 2 * c.n_kv_heads * hd + d
    n_up = 2 if c.mlp == "swiglu" else 1
    mlp = (n_up + 1) * d * c.d_ff + ((n_up * c.d_ff + d) if c.bias else 0)
    norm_one = 2 * d if c.norm == "layernorm" else d
    out = {
        "embedding": c.vocab * d,
        "position": c.max_ctx * d if c.learned_pos else 0,
        "attention": c.n_layers * attn,
        "mlp": c.n_layers * mlp,
        "norms": (2 * c.n_layers + 1) * norm_one,
        "lm_head": 0 if c.tied else c.vocab * d,
    }
    out["total"] = sum(out.values())
    return out


def count_params(c):
    return param_breakdown(c)["total"]


def matmul_params(c):
    """
    参与矩阵乘法的参数：注意力 + MLP 的权重矩阵，加上输出层（共享时也要做一次 d × V 的乘法）。
    输入嵌入只是查表，不做乘法，不计入；偏置、归一化参数很少，也忽略。用于估算 FLOPs。
    """
    d, hd = c.d_model, c.head_dim
    attn = d * c.n_heads * hd * 2 + d * c.n_kv_heads * hd * 2
    mlp = (3 if c.mlp == "swiglu" else 2) * d * c.d_ff
    return c.n_layers * (attn + mlp) + c.vocab * d


# 训练时每个参数要占多少字节（不含激活值）
RECIPES = {
    # 名字: (权重, 梯度, FP32 主权重, Adam m, Adam v) 的字节数
    "fp32 + Adam": (4, 4, 0, 4, 4),
    "混合精度 (bf16) + Adam": (2, 2, 4, 4, 4),
    "fp32 + SGD（无动量）": (4, 4, 0, 0, 0),
    "只推理 (bf16)": (2, 0, 0, 0, 0),
}


def training_memory(c, recipe="混合精度 (bf16) + Adam"):
    """
    权重、梯度、优化器状态的显存（不含激活值）。

    【返回】dict：weights、grads、master、adam_m、adam_v、total（字节）、bytes_per_param。
    【直觉】混合精度 + Adam 是 2 + 2 + 4 + 4 + 4 = 16 字节 / 参数：和纯 FP32 + Adam 一样多！
            混合精度省的是激活值的显存和计算时间，不省"参数相关"的显存。
    """
    n = count_params(c)
    w, g, m32, am, av = RECIPES[recipe]
    out = {"weights": n * w, "grads": n * g, "master": n * m32, "adam_m": n * am, "adam_v": n * av}
    out["total"] = sum(out.values())
    out["bytes_per_param"] = w + g + m32 + am + av
    return out


def activation_estimate(c, batch, seq, flash=False):
    """
    激活值显存的量级估计（Korthikanti et al., 2022：16 位激活、不重算、含 dropout 掩码时，
    每层约 s·b·h·(34 + 5·a·s/h) 字节，s = 序列长度，b = 批大小，h = 隐藏维度，a = 头数）。

    【参数】flash=True 表示用 FlashAttention 一类的方法，不存 T × T 的注意力矩阵（去掉 5·a·s/h 那一项）。
    【返回】字节数（全部层）。
    【边界】只是量级：Llama 没有 dropout、MLP 更宽，常数会不同；精确值请实测。
    """
    s, b, h, a = seq, batch, c.d_model, c.n_heads
    per_layer = s * b * h * 34 + (0 if flash else 5 * a * s * s * b)
    return per_layer * c.n_layers


def kv_cache_bytes(c, batch, seq, bytes_per=2):
    """
    推理时 KV 缓存的字节数 = 2（K 和 V）× 层数 × KV 头数 × 每头维度 × 序列长度 × 批大小 × 每个数的字节数。

    【直觉】生成第 t 个 token 时，要用到之前所有 token 在每一层的 K、V；把它们存下来就不必重算（第 19 讲）。
            GQA 让多个查询头共用一组 K/V，KV 缓存按比例缩小。
    """
    return 2 * c.n_layers * c.n_kv_heads * c.head_dim * seq * batch * bytes_per


def train_flops(c, tokens, seq):
    """
    训练的浮点运算次数 ≈ 6 × (参与矩阵乘法的参数) × token 数 + 注意力项 12 × L × seq × d × token 数。

    【为什么是 6】前向时每个参数参与一次乘法和一次加法（2 FLOPs / 参数 / token）；
                 反向要算"对输入的梯度"和"对权重的梯度"两次矩阵乘法（4 FLOPs）。合计 6。
    【注意力项】QKᵀ 和 注意力 × V 的计算量和序列长度成正比，不属于"参数"，单独加上；
                 它和主项之比约为 seq / (6d)，序列不太长时可以忽略。
    """
    return 6 * matmul_params(c) * tokens + 12 * c.n_layers * seq * c.d_model * tokens


# 公开规格（稠密 BF16 张量核心算力；不含稀疏加速）
GPUS = {
    "A100-80GB (SXM)": {"mem_gib": 80, "bf16_tflops": 312, "hbm_tb_s": 2.0},
    "H100-80GB (SXM)": {"mem_gib": 80, "bf16_tflops": 989, "hbm_tb_s": 3.35},
}


def gpu_hours(flops, gpu, mfu):
    """
    需要多少 GPU 小时 = FLOPs ÷ (单卡峰值 × MFU) ÷ 3600。

    【参数】gpu：GPUS 里的名字；mfu：模型 FLOPs 利用率（实际达到的算力 / 峰值），0～1。
    """
    return flops / (GPUS[gpu]["bf16_tflops"] * 1e12 * mfu) / 3600


def fmt_bytes(n):
    """把字节数格式化成 KiB / MiB / GiB / TiB。"""
    for unit, size in [("TiB", 1024 ** 4), ("GiB", GIB), ("MiB", 1024 ** 2), ("KiB", 1024)]:
        if n >= size:
            return f"{n / size:.1f} {unit}"
    return f"{n:.0f} B"


def fmt_count(n):
    """把参数个数格式化成 M / B。"""
    if n >= 1e9:
        return f"{n / 1e9:.2f}B"
    if n >= 1e6:
        return f"{n / 1e6:.1f}M"
    return f"{n:,}"


def small(name="tiny", **overrides):
    """一个能在 CPU 上跑的小配置，字段可以覆盖。"""
    base = Config(name, vocab=1000, d_model=256, n_layers=4, n_heads=8, n_kv_heads=8, d_ff=1024,
                  mlp="gelu", norm="layernorm", bias=True, tied=True, max_ctx=512, learned_pos=True)
    return replace(base, **overrides)