"""
=============================================================================
第 2 讲 · 代码 3：einops —— 用"维度名字"操作张量
=============================================================================

对应讲义：lectures/lecture02_pytorch_resource_accounting.md 第 4 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
写 Transformer 时，最容易出 bug 的地方不是数学，而是**维度**：
    x.view(b, h, s, d).transpose(1, 2) 到底对不对？dim=-2 指的是哪一维？
更可怕的是，很多维度错误**不会报错**：形状恰好对得上，结果却是错的。

einops 的思路：**给每个维度起名字**，用一个字符串描述"输入长什么样 → 输出长什么样"：
    rearrange(x, "batch seq (heads dim) -> batch heads seq dim", heads=8)
代码即文档，读的人一眼就知道发生了什么，写错了也更容易发现。

CS336 的作业鼓励使用 einops；后面所有讲次的代码也会大量使用它。

-----------------------------------------------------------------------------
【整体思路】
-----------------------------------------------------------------------------
    第 1 部分：einsum —— 广义的矩阵乘法（"输出里没出现的维度就求和"）
    第 2 部分：rearrange —— 转置、拆分、合并维度
    第 3 部分：reduce —— 沿某些维度求和/平均/最大值
    第 4 部分：repeat —— 复制（例如 GQA 中把 KV 头复制给多个查询头）
    第 5 部分：一个"不报错的 bug"：view 和 transpose 用错顺序
    第 6 部分：给函数写上形状注释（jaxtyping 风格）

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 03_einops.py

需要安装 einops：pip install einops
"""

import sys

import torch

try:
    from einops import einsum, rearrange, reduce, repeat
except ImportError:
    print("本文件需要 einops 库，请先运行：pip install einops")
    sys.exit(1)     # 【Python 小知识】sys.exit(1) 立即结束程序，1 表示"异常退出"


def print_header(title):
    """打印一个带分隔线的标题。"""
    print()
    print("=" * 70)
    print(f"  {title}")
    print("=" * 70)


# =============================================================================
# 第 1 部分：einsum
# =============================================================================
# 【思路】
#   einsum（爱因斯坦求和）用一个字符串描述运算：
#       einsum(x, y, "seq1 hidden, hidden seq2 -> seq1 seq2")
#   规则只有一条：**在输入里出现、但在输出里没出现的维度，会被求和掉**。
#   上面这个例子中 hidden 没出现在输出里 → 对 hidden 求和 → 这就是矩阵乘法。
#
#   einops 版的 einsum 允许用完整的英文单词作维度名（torch.einsum 只允许单个字母），可读性更好。


def demo_einsum():
    """演示 1：用 einsum 表达各种乘法。"""
    print_header("演示 1：einsum —— 输出里没有的维度就求和")
    torch.manual_seed(0)

    # (a) 矩阵乘法
    x = torch.randn(3, 4)
    y = torch.randn(4, 5)
    z = einsum(x, y, "seq1 hidden, hidden seq2 -> seq1 seq2")
    assert torch.allclose(z, x @ y, atol=1e-6)
    print(f"  矩阵乘法：(3,4) × (4,5) → {tuple(z.shape)}   ✓ 与 x @ y 一致")

    # (b) 带 batch 的矩阵乘法："..." 表示"前面任意多个维度，原样保留"
    xb = torch.randn(2, 3, 4)
    zb = einsum(xb, y, "... hidden, hidden out -> ... out")
    assert torch.allclose(zb, xb @ y, atol=1e-6)
    print(f"  批量矩阵乘法：(2,3,4) × (4,5) → {tuple(zb.shape)}   ✓")

    # (c) 注意力分数：Q 和 K 的点积。每个 query 位置与每个 key 位置之间一个分数
    B, H, S, D = 2, 4, 6, 8
    q = torch.randn(B, H, S, D)
    k = torch.randn(B, H, S, D)
    scores = einsum(q, k, "batch heads q_pos dim, batch heads k_pos dim -> batch heads q_pos k_pos")
    assert torch.allclose(scores, q @ k.transpose(-2, -1), atol=1e-5)
    print(f"  注意力分数：Q{tuple(q.shape)} · K{tuple(k.shape)} → {tuple(scores.shape)}   ✓ 与 q @ k.transpose(-2,-1) 一致")
    print("    对比一下两种写法，哪个更容易读懂？")
    print("      q @ k.transpose(-2, -1)")
    print("      einsum(q, k, 'batch heads q_pos dim, batch heads k_pos dim -> batch heads q_pos k_pos')")

    # (d) 外积：没有维度被求和
    a = torch.arange(3.0)
    b = torch.arange(4.0)
    outer = einsum(a, b, "i, j -> i j")
    print(f"  外积：(3,) × (4,) → {tuple(outer.shape)}（没有被求和的维度）")

    # (e) 逐元素相乘再对最后一维求和 = 批量点积
    u = torch.randn(5, 8)
    v = torch.randn(5, 8)
    dots = einsum(u, v, "n d, n d -> n")
    assert torch.allclose(dots, (u * v).sum(-1), atol=1e-6)
    print(f"  批量点积：(5,8)·(5,8) → {tuple(dots.shape)}   ✓")


# =============================================================================
# 第 2 部分：rearrange
# =============================================================================
def demo_rearrange():
    """演示 2：rearrange 拆分、合并、转置维度。"""
    print_header("演示 2：rearrange —— 拆分 / 合并 / 转置")
    B, S, H, D = 2, 5, 4, 8
    x = torch.randn(B, S, H * D)       # 线性层输出的形状：(batch, seq, heads*dim)

    # 拆出多个注意力头：括号 (heads dim) 表示"这一维其实是 heads 和 dim 合并来的"
    xh = rearrange(x, "batch seq (heads dim) -> batch heads seq dim", heads=H)
    print(f"  拆分注意力头：{tuple(x.shape)} → {tuple(xh.shape)}")

    # 合并回去：做完注意力后，要把各个头拼回一个向量
    back = rearrange(xh, "batch heads seq dim -> batch seq (heads dim)")
    assert torch.equal(back, x)
    print(f"  合并注意力头：{tuple(xh.shape)} → {tuple(back.shape)}   ✓ 与原张量完全相同")

    # 与纯 PyTorch 写法对比
    manual = x.view(B, S, H, D).transpose(1, 2)
    assert torch.equal(manual, xh)
    print("  等价的 PyTorch 写法：x.view(B, S, H, D).transpose(1, 2)")

    # 把图像切成小块（patch）：第 17 讲多模态会用到
    img = torch.randn(1, 3, 32, 32)    # (batch, 通道, 高, 宽)
    patches = rearrange(img, "b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1=8, p2=8)
    print(f"  图像切块：{tuple(img.shape)} → {tuple(patches.shape)}（16 个 8×8 的块，每块 192 个数）")

    # 【边界】维度大小不能整除时会报错，而且报错信息很清楚
    try:
        rearrange(x, "batch seq (heads dim) -> batch heads seq dim", heads=5)
    except Exception as e:
        print(f"  【边界】32 维拆成 5 个头 → 报错：{type(e).__name__}")


# =============================================================================
# 第 3 部分：reduce
# =============================================================================
def demo_reduce():
    """演示 3：reduce 沿指定维度做归约。"""
    print_header("演示 3：reduce —— 求和 / 平均 / 最大值")
    x = torch.randn(2, 5, 8)           # (batch, seq, hidden)

    m = reduce(x, "batch seq hidden -> batch hidden", "mean")
    assert torch.allclose(m, x.mean(dim=1))
    print(f"  对序列维求平均：{tuple(x.shape)} → {tuple(m.shape)}   ✓ 与 x.mean(dim=1) 一致")

    s = reduce(x, "... hidden -> ...", "sum")
    assert torch.allclose(s, x.sum(dim=-1))
    print(f"  对最后一维求和：{tuple(x.shape)} → {tuple(s.shape)}   ✓ 与 x.sum(dim=-1) 一致")

    # 分组池化：把 hidden 维看成 (hidden g)，即每 g=4 个相邻元素一组，每组取最大值。
    # 【边界】被分组的维度必须能整除 g；这里 hidden=8 可以被 4 整除。
    p = reduce(x, "batch seq (hidden g) -> batch seq hidden", "max", g=4)
    print(f"  隐藏维每 4 个取最大值：{tuple(x.shape)} → {tuple(p.shape)}")


# =============================================================================
# 第 4 部分：repeat
# =============================================================================
def demo_repeat():
    """演示 4：repeat —— 在 GQA 中复制 KV 头。"""
    print_header("演示 4：repeat —— 复制（GQA 的例子）")
    # 分组查询注意力（GQA，第 3 讲）：8 个查询头共享 2 个 KV 头，每个 KV 头被 4 个查询头使用
    B, S, D = 1, 6, 16
    kv_heads, q_heads = 2, 8
    k = torch.randn(B, kv_heads, S, D)
    k_expanded = repeat(k, "b kv s d -> b (kv g) s d", g=q_heads // kv_heads)
    print(f"  KV 头 {tuple(k.shape)} → 复制给查询头 {tuple(k_expanded.shape)}")
    assert torch.equal(k_expanded[:, 0], k_expanded[:, 3]) and torch.equal(k_expanded[:, 4], k[:, 1])
    print("  ✓ 查询头 0~3 共用 KV 头 0，查询头 4~7 共用 KV 头 1")


# =============================================================================
# 第 5 部分：一个不报错的 bug
# =============================================================================
def demo_silent_bug():
    """演示 5：形状对了，结果却错了。"""
    print_header("演示 5：最危险的 bug —— 形状对了，内容错了")
    B, S, H, D = 1, 3, 2, 4
    x = torch.arange(B * S * H * D, dtype=torch.float32).view(B, S, H * D)

    right = rearrange(x, "b s (h d) -> b h s d", h=H)
    wrong = x.view(B, H, S, D)         # 想当然地直接 view 成目标形状

    print(f"  正确结果形状 {tuple(right.shape)}，错误结果形状 {tuple(wrong.shape)}   ← 形状完全一样！")
    print(f"  内容是否相同：{torch.equal(right, wrong)}")
    print(f"  正确：第 0 个头、第 1 个位置 = {right[0, 0, 1].tolist()}")
    print(f"  错误：第 0 个头、第 1 个位置 = {wrong[0, 0, 1].tolist()}")
    print("""
  x.view(B, H, S, D) 只是把内存里的数"重新按这个形状读一遍"，
  但内存里的顺序是 (seq, heads, dim)，不是 (heads, seq, dim)。
  正确做法是先 view 成 (B, S, H, D) 再交换 seq 和 heads 两维。
  这种错误不会报错，模型照样能训练，只是效果莫名其妙地变差 —— 非常难排查。
  用 einops 写出维度名字，这类错误一眼就能看出来。""")


# =============================================================================
# 第 6 部分：形状注释
# =============================================================================
# 【思路】
#   CS336 推荐用 jaxtyping 库给张量写上形状注释，例如：
#       def attention(q: Float[Tensor, "batch heads seq d"], ...) -> Float[Tensor, "batch heads seq d"]
#   这些注释默认不会被检查（除非配合运行时检查工具），但它们是极好的文档。
#   为了不增加依赖，本课程的代码在文档字符串里用同样的风格写形状，例如下面这个函数。


def scaled_dot_product_scores(q, k):
    """
    计算缩放后的注意力分数。

    【形状】
        q      : (batch, heads, q_len, d)
        k      : (batch, heads, k_len, d)
        返回值 : (batch, heads, q_len, k_len)

    【思路】点积除以 sqrt(d)（为什么要除，第 3 讲会讲）。
    """
    d = q.shape[-1]
    return einsum(q, k, "b h q d, b h k d -> b h q k") / d ** 0.5


def demo_shape_annotations():
    """演示 6：带形状注释的函数。"""
    print_header("演示 6：形状注释")
    q = torch.randn(2, 4, 5, 16)
    k = torch.randn(2, 4, 7, 16)
    s = scaled_dot_product_scores(q, k)
    print(f"  q{tuple(q.shape)}, k{tuple(k.shape)} → 分数 {tuple(s.shape)}（q_len=5, k_len=7 可以不同）")


def main():
    demo_einsum()
    demo_rearrange()
    demo_reduce()
    demo_repeat()
    demo_silent_bug()
    demo_shape_annotations()
    print("\n全部演示运行完毕。")


if __name__ == "__main__":
    main()