"""
=============================================================================
CS349E · 第 4 讲共享模块：shard_lab.py —— 在虚拟 GPU 上切分矩阵、做张量并行，并给每一次通信记账
=============================================================================

【这个模块要解决什么问题】
    一个模型放不进一块芯片（或一块芯片太慢），就要把矩阵切到很多块芯片上。切开以后有两件事要搞清楚：
        1. 数学上还对不对：每块芯片只有一部分数据，怎样拼出和单卡完全一样的结果？
        2. 代价是什么：为了拼出结果，芯片之间要传多少字节、花多少时间？
    本模块在一台电脑上模拟 n 块"虚拟 GPU"：每块 GPU 持有真实的 numpy 数组分片，集合通信用前置讲 P5 的
    comm_sim（真实搬运数据 + 逐消息计时）。在此基础上实现：
        - 切分 / 合并（shard / unshard）与三种集合通信的包装（all-gather、reduce-scatter、all-reduce），
          约定"设备 d 持有第 d 块"；
        - How to Scale Your Model 第 3 章"切分矩阵乘法的四种情形"；
        - 第 2 讲 Transformer 的 Megatron 张量并行前向（按头切注意力、按 F 切 MLP、按词表切 lm_head），
          可以和单卡结果逐个比较；
        - 解析的代价公式：张量并行 decode 一步的时间（计算 + 访存 + 通信）、Pope et al. 的分区通信量。

【切分记号（How to Scale Your Model / Pope et al.）】
    A[I_X, J] 表示 A 的第一维 I 被切到网格轴 X 上（每块设备拿 I / |X| 行），J 没被切（每块设备都有完整的 J）。
    没写下标的维度是"复制"的。本模块只用一条网格轴（n 块设备挂在一个 NVSwitch 上）。

【运行】python3 shard_lab.py   # 自检：集合通信结果正确；张量并行的输出 = 单卡输出
【依赖】numpy；../prereq5/comm_sim、../lecture02/transformer_lab、../lecture01 的 hw_specs / llm_math。
"""

import math
import os
import sys

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
for sub in ("lecture01", "lecture02", "prereq5"):
    sys.path.insert(0, os.path.join(HERE, "..", sub))

from comm_sim import (ALLREDUCE_ALGOS, NVLINK4, Comm, allgather_ring, reduce_scatter_ring, switch,  # noqa: E402
                      t_formula)
from transformer_lab import KVCache, apply_rope, attention_core, rms_norm, rope_angles, silu  # noqa: E402


# ---------------------------------------------------------------------------
# 切分与集合通信（"设备 d 持有第 d 块"）
# ---------------------------------------------------------------------------
def make_comm(n, link=NVLINK4, byte_scale=1.0):
    """n 块设备挂在一个无阻塞交换机上（NVSwitch 一类），返回 comm_sim.Comm。"""
    return Comm(switch(n, link), byte_scale=byte_scale)


def shard(x, axis, n):
    """
    把数组 x 沿 axis 切成 n 块，第 d 块给设备 d。

    【边界】axis 的长度必须能被 n 整除（真实系统会补齐 padding；这里直接报错，避免悄悄算错）。
    【Python 小知识】np.split(x, n, axis) 等分成 n 块，返回列表。
    """
    if x.shape[axis] % n:
        raise ValueError(f"维度 {axis} 的长度 {x.shape[axis]} 不能被 {n} 整除")
    return np.split(x, n, axis=axis)


def unshard(shards, axis):
    """把各设备的分片按设备顺序拼回来（只在检查结果时用；真实系统里没有"上帝视角"）。"""
    return np.concatenate(shards, axis=axis)


def all_gather(comm, shards, axis):
    """
    all-gather：每块设备都拿到完整的数组（沿 axis 拼接所有分片）。用 comm_sim 的环形算法搬运并计时。

    【思路】comm_sim 的环形算法处理一维数组、按块号顺序拼接。先把 axis 挪到第 0 维再拉平：
        "沿第 0 维拼接"和"把拉平后的一维数组首尾相连"是同一件事。
    【返回】列表：每个设备一个完整数组
    【Python 小知识】np.moveaxis(x, axis, 0) 把第 axis 维挪到最前面（只改视图，不复制数据）。
    """
    n = len(shards)
    moved = [np.moveaxis(s, axis, 0) for s in shards]
    flat = [m.ravel() for m in moved]
    full = allgather_ring(comm, flat, {d: d for d in range(n)}, label="all-gather")
    shape = (moved[0].shape[0] * n,) + moved[0].shape[1:]
    return [np.moveaxis(f.reshape(shape), 0, axis) for f in full]


def reduce_scatter(comm, arrays, axis):
    """
    reduce-scatter：各设备持有同样形状的"部分和"，结果是"和"沿 axis 切开，设备 d 拿第 d 块。

    【思路】comm_sim 的环形 reduce-scatter 结束时，设备 r 持有第 (r+1) mod n 块。所以先把每个设备的块
        "转一格"排列（第 c 块放规范块 (c−1) mod n），这样设备 r 最后拿到的正好是规范的第 r 块。
    【返回】列表：设备 d 的结果分片
    """
    n = len(arrays)
    moved = [np.moveaxis(a, axis, 0) for a in arrays]
    if moved[0].shape[0] % n:
        raise ValueError("reduce-scatter 的维度必须能被设备数整除")
    blocks = [np.split(m, n, axis=0) for m in moved]
    rotated = [np.concatenate([b[(c - 1) % n].ravel() for c in range(n)]) for b in blocks]
    pieces, owner = reduce_scatter_ring(comm, rotated, label="reduce-scatter")
    shape = blocks[0][0].shape
    out = []
    for d in range(n):
        assert (owner[d] - 1) % n == d
        out.append(np.moveaxis(pieces[d].reshape(shape), 0, axis))
    return out


def all_reduce(comm, arrays, algo="ring"):
    """
    all-reduce：各设备的数组求和，每个设备都拿到和。algo 取 comm_sim.ALLREDUCE_ALGOS 的键（ring、rd、innet……）。

    【返回】列表：每个设备一个（相同的）和
    """
    shape = arrays[0].shape
    fn = ALLREDUCE_ALGOS[algo][1]
    out = fn(comm, [a.ravel() for a in arrays])
    return [o.reshape(shape) for o in out]


# ---------------------------------------------------------------------------
# 切分矩阵乘法的四种情形（How to Scale Your Model 第 3 章）
# ---------------------------------------------------------------------------
def matmul_case1(A_sh, B):
    """
    情形 1：A[I_X, J] · B[J, K] → C[I_X, K]。被求和的维度 J 没被切：每块设备用自己的行做本地乘法，不需要通信。
    【返回】C 的分片列表
    """
    return [a @ B for a in A_sh]


def matmul_case2(comm, A_sh, B):
    """
    情形 2：A[I, J_X] · B[J, K]（B 完整复制）。被求和的维度只在 A 上被切：先 all-gather A，再本地乘。
    【返回】每块设备一个完整的 C[I, K]（结果是复制的，每块设备都做了全部运算）
    【另一种做法】B 是复制的，每块设备也可以取 B[J_X, K] 那一片做部分和，再 all-reduce（见 matmul_case3）——
        运算被分摊了，但通信的是 [I, K] 而不是 [I, J]。哪种划算取决于 J 和 K 谁大。
    """
    A_full = all_gather(comm, A_sh, axis=1)
    return [a @ B for a in A_full]


def matmul_case3(comm, A_sh, B_sh, scatter_axis=None, algo="ring"):
    """
    情形 3：A[I, J_X] · B[J_X, K]。两边被求和的维度切法一致：每块设备算一个部分和 C[I, K]{U_X}
    （"还没加起来"），再 all-reduce 得到完整的 C；或者 reduce-scatter 得到 C[I, K_X]（scatter_axis=1）。
    【返回】all-reduce：每块设备一个完整 C；reduce-scatter：C 的分片列表
    """
    partial = [a @ b for a, b in zip(A_sh, B_sh)]
    if scatter_axis is None:
        return all_reduce(comm, partial, algo)
    return reduce_scatter(comm, partial, scatter_axis)


def matmul_case4(comm, A_sh, B_sh):
    """
    情形 4：A[I_X, J] · B[J, K_X]。两个"不被求和"的维度切在同一条轴上：设备 d 只有 I 的第 d 段和 K 的第 d 段，
    只能算出 C 的对角块——C[I_X, K_X] 不是一个合法的切法（会漏掉非对角块）。
    做法：先 all-gather 其中一个（这里 all-gather B），变成情形 1：C[I_X, K]。
    【返回】C 的分片列表（沿 I 切）
    """
    B_full = all_gather(comm, B_sh, axis=1)
    return [a @ b for a, b in zip(A_sh, B_full)]


# ---------------------------------------------------------------------------
# Megatron 张量并行（第 2 讲的 Transformer）
# ---------------------------------------------------------------------------
def kv_group_of_device(cfg, n, d):
    """
    张量并行时设备 d 负责的 KV 组（列表）。

    【思路】query 头均分给 n 块设备（每块 H/n 个）。设备 d 的第一个头是 d·H/n，它属于第 (d·H/n) // R 组（R = H/G）。
        n ≤ G：每块设备分到 G/n 个完整的组，KV 不重复；
        n > G：每块设备只有一个组的一部分 query 头，但那一整组 K/V 都得有——**KV 被复制** n/G 份。
    """
    H, G = cfg.n_heads, cfg.n_kv_heads
    if H % n:
        raise ValueError(f"query 头数 {H} 不能被设备数 {n} 整除")
    R = H // G
    hpd = H // n
    first, last = d * hpd, (d + 1) * hpd - 1
    return list(range(first // R, last // R + 1))


def shard_params_tp(params, cfg, n):
    """
    把第 2 讲模型的参数按 Megatron 的方式切给 n 块设备。

    【切法】
        Wq：按列（query 头）切，每块 H/n 个头；Wk、Wv：取这块设备的 query 头所属的 KV 组的列（n > G 时复制）；
        Wo：按行切（和 Wq 的头对应）——"先列后行"，中间不需要通信；
        W_gate、W_up：按列切（F 维），W_down：按行切（F 维）——同样"先列后行"；
        lm_head：按列切（词表维），最后 all-gather logits；
        词嵌入、RMSNorm 的缩放系数：每块设备一份完整的（复制）。
    【返回】列表：每块设备一个参数字典（结构和单卡相同，只是矩阵变窄了），外加 "kv_groups"
    """
    H, G, dh, d_model = cfg.n_heads, cfg.n_kv_heads, cfg.d_head, cfg.d_model
    hpd = H // n
    if cfg.d_ff % n or cfg.vocab % n:
        raise ValueError("d_ff 和 vocab 必须能被设备数整除")
    out = []
    for d in range(n):
        groups = kv_group_of_device(cfg, n, d)
        q_cols = slice(d * hpd * dh, (d + 1) * hpd * dh)
        kv_cols = np.concatenate([np.arange(g * dh, (g + 1) * dh) for g in groups])
        f_cols = slice(d * cfg.d_ff // n, (d + 1) * cfg.d_ff // n)
        v_cols = slice(d * cfg.vocab // n, (d + 1) * cfg.vocab // n)
        dp = {"embed": params["embed"], "final_norm": params["final_norm"], "kv_groups": groups, "layers": []}
        head = params["lm_head"] if params["lm_head"] is not None else params["embed"].T
        dp["lm_head"] = head[:, v_cols]
        for lp in params["layers"]:
            dp["layers"].append({
                "attn_norm": lp["attn_norm"], "mlp_norm": lp["mlp_norm"],
                "wq": lp["wq"][:, q_cols], "wk": lp["wk"][:, kv_cols], "wv": lp["wv"][:, kv_cols],
                "wo": lp["wo"][q_cols, :],
                "w_gate": lp["w_gate"][:, f_cols], "w_up": lp["w_up"][:, f_cols], "w_down": lp["w_down"][f_cols, :],
            })
        out.append(dp)
    return out


def tp_forward(dev_params, cfg, tokens, comm, caches=None, algo="ring", gather_logits=True):
    """
    张量并行的一次前向：n 块设备各算自己那部分头 / 那部分 F，每层两次 all-reduce。

    【参数】dev_params：shard_params_tp 的结果；tokens [B, T]；comm：comm_sim.Comm（记录通信）；
        caches：每块设备一个 KVCache（只存自己的 KV 组），None 表示不用 cache；algo：all-reduce 算法；
        gather_logits：True 时 all-gather 各设备的词表分片，返回完整 logits
    【返回】完整 logits [B, T, vocab]（每块设备都一样，返回设备 0 的）
    【思路（每层）】
        x 在每块设备上是一样的（复制）。h = RMSNorm(x)（每块设备各自算一遍，重复但便宜）；
        注意力：设备 d 用自己的 Wq / Wk / Wv 列算自己那几个头，乘自己的 Wo 行 → 部分和 → all-reduce → 加到 x；
        MLP：设备 d 用自己的 W_gate / W_up 列算 F/n 个中间值，乘自己的 W_down 行 → 部分和 → all-reduce → 加到 x。
        "先按列切、再按行切"让两个矩阵乘法之间不需要通信，整个子层只在最后 all-reduce 一次（Megatron 的技巧）。
    """
    n = len(dev_params)
    B, T = tokens.shape
    dh = cfg.d_head
    pos0 = caches[0].length if caches is not None else 0
    cos, sin = rope_angles(np.arange(pos0, pos0 + T), dh, cfg.rope_theta)
    xs = [p["embed"][tokens] for p in dev_params]
    for li in range(cfg.n_layers):
        partial = []
        for d, p in enumerate(dev_params):
            lp = p["layers"][li]
            h = rms_norm(xs[d], lp["attn_norm"], cfg.norm_eps)
            nh = lp["wq"].shape[1] // dh
            ng = lp["wk"].shape[1] // dh
            q = apply_rope((h @ lp["wq"]).reshape(B, T, nh, dh), cos, sin)
            k = apply_rope((h @ lp["wk"]).reshape(B, T, ng, dh), cos, sin)
            v = (h @ lp["wv"]).reshape(B, T, ng, dh)
            if caches is not None:
                caches[d].write(li, k, v)
                K, V = caches[d].read(li, pos0 + T)
            else:
                K, V = k, v
            if nh % ng:
                # n > G 时，一块设备的 query 头只占一个 KV 组的一部分：attention_core 需要 nh 是 ng 的倍数
                raise ValueError("每块设备的 query 头数必须是 KV 组数的倍数")
            a = attention_core(q, K, V, pos0)
            partial.append(a @ lp["wo"])
        summed = all_reduce(comm, partial, algo)
        xs = [x + s for x, s in zip(xs, summed)]
        partial = []
        for d, p in enumerate(dev_params):
            lp = p["layers"][li]
            h = rms_norm(xs[d], lp["mlp_norm"], cfg.norm_eps)
            m = silu(h @ lp["w_gate"]) * (h @ lp["w_up"])
            partial.append(m @ lp["w_down"])
        summed = all_reduce(comm, partial, algo)
        xs = [x + s for x, s in zip(xs, summed)]
    if caches is not None:
        for c in caches:
            c.length += T
    local = [rms_norm(x, p["final_norm"], cfg.norm_eps) @ p["lm_head"] for x, p in zip(xs, dev_params)]
    if not gather_logits:
        return local
    return all_gather(comm, local, axis=2)[0]


def tp_caches(cfg, dev_params, batch):
    """为每块设备建一个只存自己那几个 KV 组的 KV cache。"""
    from dataclasses import replace
    out = []
    for p in dev_params:
        c = replace(cfg, n_kv_heads=len(p["kv_groups"]), n_heads=len(p["kv_groups"]))
        out.append(KVCache(c, batch, max_len=cfg.max_seq))
    return out


# ---------------------------------------------------------------------------
# 解析代价模型
# ---------------------------------------------------------------------------
def tp_decode_step(model, chip, link, n, batch, ctx, w_bytes=2.0, kv_bytes=2.0, act_bytes=2.0, algo="ring",
                   compute_prec="bf16"):
    """
    张量并行 decode 一步的 roofline 时间：每块芯片的计算 / 访存（理想地除以 n）+ 每层两次 all-reduce。

    【参数】model：llm_math.ModelConfig；chip：hw_specs.Chip；link：comm_sim.LinkSpec（α、β）；n：张量并行度；
        batch、ctx：批大小、上下文长度；w_bytes / kv_bytes / act_bytes：权重 / KV / 激活每个数的字节；
        algo：all-reduce 算法（"ring"、"rd"……，用 comm_sim.t_formula 的公式）
    【返回】字典：t_math、t_mem、t_comm、t_step（= max(t_math, t_mem) + t_comm，通信不和计算重叠，偏保守）、
        kv_per_chip（每块芯片存多少 KV 字节 / token / 序列）、n_allreduce
    【KV 的复制】n > G（KV 组数）时，每块芯片也要存一整组 KV：每块芯片的 KV = 总 KV × max(1, n/G) / n。
    【边界】n = 1 时没有通信。
    """
    w = model.matmul_params() * w_bytes / n
    kv_tok = model.kv_bytes_per_token(kv_bytes) * max(1.0, n / model.n_kv_heads) / n
    kv = batch * ctx * kv_tok
    flops = batch * model.flops_per_token(ctx) / n
    t_math = flops / chip.peak_flops(compute_prec)
    t_mem = (w + kv) / chip.mem_bw
    n_ar = 2 * model.n_layers
    msg = batch * model.d_model * act_bytes
    t_comm = 0.0 if n == 1 else n_ar * t_formula(algo, n, msg, link.alpha, link.bw)
    return {"t_math": t_math, "t_mem": t_mem, "t_comm": t_comm, "t_step": max(t_math, t_mem) + t_comm,
            "kv_per_chip_per_token": kv_tok, "n_allreduce": n_ar, "msg_bytes": msg}


def pope_ffn_comm(variant, tokens, E, F, n, bw, bytes_per_elem=2.0):
    """
    Pope et al. §3.2 的前馈层通信时间公式（每层，单位：秒）。

        1D weight-stationary：2·BL·E / 带宽                      （和芯片数无关）
        2D weight-stationary：8·BL·E / (√n · 带宽)               （F = 4E 时最优的 X、YZ 切法）
        weight-gathered     ：2·E·F·N / (n · 带宽) + 2·BL·E / (N · 带宽)
                               （附录 A.2.2：权重在 N 块芯片之间 all-gather，激活在其余方向聚合；
                                 N = √(BL·n / F) 时两项相等、总和最小，= 4·E·√(BL·F) / (√n · 带宽)）
    【参数】variant："1d"、"2d" 或 "gathered"；tokens = B·L（这一步的 token 数）；E、F：模型维、前馈维；
        n：芯片数；bw：每块芯片的网络带宽（字节/秒）；bytes_per_elem：每个数的字节
    【返回】秒；variant == "gathered" 时返回 (秒, 实际用的 N)
    【边界】公式忽略了延迟 α，只看带宽项；N 只能在 1 和 n 之间（最优值超出时取边界）。
    """
    if variant == "1d":
        return 2 * tokens * E * bytes_per_elem / bw
    if variant == "2d":
        return 8 * tokens * E * bytes_per_elem / (math.sqrt(n) * bw)
    if variant == "gathered":
        N = min(max(math.sqrt(tokens * n / F), 1.0), float(n))
        return (2 * E * F * N / n + 2 * tokens * E / N) * bytes_per_elem / bw, N
    raise ValueError(variant)


# ---------------------------------------------------------------------------
# 自检
# ---------------------------------------------------------------------------
def _selftest():
    from transformer_lab import TINY, forward, init_params
    rng = np.random.default_rng(0)
    n = 4
    comm = make_comm(n)
    A = rng.standard_normal((8, 12))
    B = rng.standard_normal((12, 16))
    C = A @ B
    assert np.allclose(unshard(matmul_case1(shard(A, 0, n), B), 0), C)
    assert all(np.allclose(c, C) for c in matmul_case2(comm, shard(A, 1, n), B))
    assert all(np.allclose(c, C) for c in matmul_case3(comm, shard(A, 1, n), shard(B, 0, n)))
    assert np.allclose(unshard(matmul_case3(comm, shard(A, 1, n), shard(B, 0, n), scatter_axis=1), 1), C)
    assert np.allclose(unshard(matmul_case4(comm, shard(A, 0, n), shard(B, 1, n)), 0), C)
    cfg = TINY
    p = init_params(cfg, seed=0)
    toks = rng.integers(0, cfg.vocab, size=(2, 9))
    ref = forward(p, cfg, toks)
    for n in (1, 2, 4, 8):
        out = tp_forward(shard_params_tp(p, cfg, n), cfg, toks, make_comm(n))
        assert np.max(np.abs(out - ref)) < 1e-10, n
    print("shard_lab 自检通过：四种切分矩阵乘法 = numpy；张量并行（1/2/4/8 块设备）的 logits = 单卡")


if __name__ == "__main__":
    _selftest()