"""Task 38 算法级测试

用纯 torch 模拟 K1/K2/K3 的逻辑, 跟平台给的 reference 比较, 验证算法正确性。
Triton 编译/运行 只能在有 GPU + Triton 的环境验证 (本机没 GPU)。
"""
import torch
import sys
sys.path.insert(0, '.')


# ============================================================
# 平台给的 reference (照抄, 完全一样)
# ============================================================
def reference(logits, k, n_shared_experts, route_scale, global_scale, bias):
    M, G = logits.shape
    N = G - n_shared_experts
    S = n_shared_experts

    logits_f = logits.float()
    routed_logits = logits_f[:, :N]
    sel = torch.sigmoid(routed_logits) + bias.float()[None, :]

    _, idx = torch.topk(sel, k, dim=-1)
    routed_vals = torch.gather(routed_logits, 1, idx)
    shared_vals = logits_f[:, N:N + S]

    active = torch.cat([routed_vals, shared_vals], dim=-1)
    probs = torch.sigmoid(active)
    weights = probs / probs.sum(dim=-1, keepdim=True)
    weights = weights * route_scale * global_scale.float()

    routed_w = weights[:, :k].to(logits.dtype)
    shared_w = weights[:, k:].to(logits.dtype)
    indices = idx.to(torch.int32)
    return routed_w, indices, shared_w


# ============================================================
# 算法模拟: 分 3 步, 跟 Triton 写法一一对应
# ============================================================
def algo_sigmoid_bias(routed, bias):
    """模拟 K1: sel = sigmoid(routed) + bias"""
    return torch.sigmoid(routed.float()) + bias.float()[None, :]


def algo_topk(sel, k):
    """模拟 K2: topk 选 indices"""
    _, idx = torch.topk(sel, k, dim=-1)
    return idx.to(torch.int32)


def algo_gate_finalize(routed, shared, indices, k, route_scale, global_scale, output_dtype):
    """模拟 K3: gather + cat + sigmoid + renorm + scale + split"""
    # gather 原始 routed
    routed_vals = torch.gather(routed.float(), 1, indices.long())
    shared_vals = shared.float()

    # cat
    active = torch.cat([routed_vals, shared_vals], dim=-1)
    # sigmoid + renorm
    probs = torch.sigmoid(active)
    weights = probs / probs.sum(dim=-1, keepdim=True)
    # scale
    weights = weights * route_scale * global_scale.float()

    # split + cast
    routed_w = weights[:, :k].to(output_dtype)
    shared_w = weights[:, k:].to(output_dtype)
    return routed_w, shared_w


def algo_full(logits, k, n_shared_experts, route_scale, global_scale, bias):
    T, G = logits.shape
    N = G - n_shared_experts
    S = n_shared_experts
    routed = logits[:, :N]
    shared = logits[:, N:]

    sel = algo_sigmoid_bias(routed, bias)
    indices = algo_topk(sel, k)
    routed_w, shared_w = algo_gate_finalize(
        routed, shared, indices, k, route_scale, global_scale, logits.dtype
    )
    return routed_w, indices, shared_w


# ============================================================
# 跑测试
# ============================================================
def test_case(T, N, S, k, dtype=torch.float16, seed=42):
    """跑一个测试用例, 比较 reference 和 algo_full"""
    torch.manual_seed(seed)
    G = N + S
    logits = torch.randn(T, G, dtype=dtype) * 2
    bias = torch.randn(N, dtype=torch.float32) * 0.1
    route_scale = 1.5
    global_scale = torch.tensor([2.0], dtype=torch.float32)

    # 平台 reference
    ref_routed_w, ref_indices, ref_shared_w = reference(
        logits, k, n_shared_experts=S, route_scale=route_scale,
        global_scale=global_scale, bias=bias
    )

    # 我们的算法
    algo_routed_w, algo_indices, algo_shared_w = algo_full(
        logits, k, n_shared_experts=S, route_scale=route_scale,
        global_scale=global_scale, bias=bias
    )

    # 比较
    idx_match = (ref_indices == algo_indices).all().item()
    routed_close = torch.allclose(ref_routed_w.float(), algo_routed_w.float(), atol=1e-3, rtol=1e-3)
    shared_close = torch.allclose(ref_shared_w.float(), algo_shared_w.float(), atol=1e-3, rtol=1e-3)

    print(f"  T={T} N={N} S={S} k={k} dtype={dtype}:")
    print(f"    indices match: {idx_match}  routed close: {routed_close}  shared close: {shared_close}")

    if not (idx_match and routed_close and shared_close):
        # 打印前几个 mismatch
        if not idx_match:
            mismatch = (ref_indices != algo_indices).nonzero()
            print(f"    First 3 idx mismatches: {mismatch[:3].tolist()}")
        if not routed_close:
            diff = (ref_routed_w.float() - algo_routed_w.float()).abs()
            print(f"    Max routed diff: {diff.max().item()}")
        if not shared_close:
            diff = (ref_shared_w.float() - algo_shared_w.float()).abs()
            print(f"    Max shared diff: {diff.max().item()}")
        return False
    return True


if __name__ == "__main__":
    print("Testing Task 38 algorithm...")
    print()

    test_cases = [
        # (T, N, S, k)
        (1, 64, 1, 8),     # tiny
        (1, 64, 1, 4),     # k < 8
        (8, 64, 1, 8),     # batch
        (32, 128, 2, 8),   # medium, group-style
        (64, 256, 1, 8),   # DeepSeek-V3-ish
        (128, 256, 1, 6),  # k=6
        (1, 8, 1, 2),      # very tiny
        (16, 32, 1, 4),    # small N
    ]

    all_pass = True
    for T, N, S, k in test_cases:
        ok = test_case(T, N, S, k)
        all_pass = all_pass and ok

    print()
    print("=" * 50)
    print(f"{'ALL PASS ✓' if all_pass else 'SOME FAILED ✗'}")
    print("=" * 50)