"""Task 55 hc_head v1 算法层测试

用 torch 模拟 v1 kernel 的逻辑, 跟平台给的 reference 比较
(bit-exact 不可能, bf16 有累积误差, 用 atol=1.5e-2, rtol=1.5e-2 平台标准)

注意: 这个测试是 torch 模拟 v1 算法, 验证"如果 v1 在真 GPU 上跑, 算法层是否对"
"""
import torch
import torch.nn.functional as F
import sys
sys.path.insert(0, '.')


# ============================================================
# 平台给的 reference (照抄)
# ============================================================
def reference(x, hc_fn, hc_scale, hc_base, norm_eps, hc_eps):
    shape, dtype = x.size(), x.dtype
    x = x.flatten(1).float()
    rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + norm_eps)
    mixes = F.linear(x, hc_fn) * rsqrt
    pre = torch.sigmoid(mixes * hc_scale + hc_base) + hc_eps
    y = torch.sum(pre.unsqueeze(-1) * x.view(shape), dim=1)
    return y.to(dtype)


# ============================================================
# v1 的 torch 模拟(模拟 K1 的 2-pass 逻辑)
# ============================================================
def v1_torch(x, hc_fn, hc_scale, hc_base, norm_eps, hc_eps):
    """v1 的 torch 模拟: 跟 K1 逻辑 bit-exact"""
    T, hc_mult, D = x.shape
    HC_DIM = hc_mult * D
    x_flat = x.view(T, HC_DIM).float()
    hc_fn_f = hc_fn.float()
    scale = hc_scale.float().item()
    bases = hc_base.float()
    
    # Pass 1: sqr_sum + mixes
    sqr_sum = x_flat.square().sum(-1)  # [T]
    rsqrt = torch.rsqrt(sqr_sum / HC_DIM + norm_eps)  # [T]
    mixes = F.linear(x_flat, hc_fn_f) * rsqrt.unsqueeze(-1)  # [T, HC_MULT]
    
    # sigmoid gate
    pre = torch.sigmoid(mixes * scale + bases) + hc_eps  # [T, HC_MULT]
    
    # Pass 2: 加权求和
    x_back = x.view(T, hc_mult, D).float()  # [T, HC_MULT, D]
    y = (pre.unsqueeze(-1) * x_back).sum(dim=1)  # [T, D]
    return y.to(x.dtype)


# ============================================================
# 跑测试
# ============================================================
def test_case(T, hc_mult, D, norm_eps=1e-6, hc_eps=1e-6, seed=42):
    """跑一个测试用例, 比较 reference 和 v1_torch"""
    torch.manual_seed(seed)
    x = torch.randn(T, hc_mult, D, dtype=torch.bfloat16) * 0.5
    hc_fn = torch.randn(hc_mult, hc_mult * D, dtype=torch.float32) * 0.02
    hc_scale = torch.tensor([2.0], dtype=torch.float32)
    hc_base = torch.randn(hc_mult, dtype=torch.float32) * 0.1

    ref_out = reference(x, hc_fn, hc_scale, hc_base, norm_eps, hc_eps)
    v1_out = v1_torch(x, hc_fn, hc_scale, hc_base, norm_eps, hc_eps)

    # 平台标准: atol=1.5e-2, rtol=1.5e-2
    match = torch.allclose(v1_out.float(), ref_out.float(), atol=1.5e-2, rtol=1.5e-2)
    max_diff = (v1_out.float() - ref_out.float()).abs().max().item()
    max_ref = ref_out.float().abs().max().item()
    rel_diff = max_diff / max(max_ref, 1e-9)

    print(f"  T={T:5d} hc_mult={hc_mult} D={D:5d}  "
          f"match={match} max_abs_diff={max_diff:.6f} max_rel_diff={rel_diff:.6f}")
    return match


if __name__ == "__main__":
    print("Testing Task 55 hc_head v1 vs reference algorithm equivalence...")
    print("(平台标准: atol=1.5e-2, rtol=1.5e-2)")
    print()

    test_cases = [
        # (T, hc_mult, D)  默认 hc_mult=4 (DeepSeek-V4), D=7168
        (1, 4, 7168),       # 1 token
        (8, 4, 7168),       # 8 token
        (32, 4, 7168),      # 32 token
        (128, 4, 7168),     # 128 token (典型 batch)
        (1, 4, 1024),       # 1 token 小 D
        (8, 4, 256),        # 8 token 更小 D
        (16, 2, 128),       # hc_mult=2, 小 D
        (1, 4, 128),        # 极小
    ]

    all_pass = True
    for T, hc_mult, D in test_cases:
        ok = test_case(T, hc_mult, D)
        all_pass = all_pass and ok

    # 再用不同 norm_eps / hc_eps 测
    print()
    print("Extra edge cases (不同 eps):")
    for eps in [1e-8, 1e-4, 1e-2]:
        ok = test_case(8, 4, 1024, norm_eps=eps, hc_eps=eps)
        all_pass = all_pass and ok

    print()
    print("=" * 70)
    print(f"{'ALL PASS ✓' if all_pass else 'SOME FAILED ✗'}")
    print("=" * 70)
    sys.exit(0 if all_pass else 1)
