"""Correctness test for all v2 operators.

Run with: python -m pytest tests/test_correctness.py -v
or:       python tests/test_correctness.py
"""
import torch


def test_gelu_and_mul_v2_correctness():
    import sys
    sys.path.insert(0, "src")
    from flag_gems.fused.gelu_and_mul import gelu_and_mul as baseline
    from flag_gems_local.fused.gelu_and_mul_v2 import gelu_and_mul as v2

    torch.manual_seed(42)
    for shape in [(1024, 1024), (4096, 4096), (2048, 11008)]:
        for dtype in [torch.float16, torch.bfloat16, torch.float32]:
            x = torch.randn(shape, dtype=dtype, device="cuda")
            y = torch.randn(shape, dtype=dtype, device="cuda")
            out_b = baseline(x, y)
            out_v = v2(x, y)
            abs_diff = (out_b - out_v).abs().max().item()
            assert abs_diff < 1e-2, (
                f"gelu_and_mul v2 diverges: shape={shape} dtype={dtype} abs_diff={abs_diff}"
            )
            print(f"gelu_and_mul  shape={shape} dtype={dtype}  abs_diff={abs_diff:.2e}  ✓")


def test_rotary_embedding_v2_correctness():
    import sys
    sys.path.insert(0, "src")
    from flag_gems.fused.rotary_embedding import apply_rotary_pos_emb as baseline
    from flag_gems_local.fused.rotary_embedding_v2 import apply_rotary_pos_emb as v2

    torch.manual_seed(42)
    for B, S, H, D in [(1, 1024, 32, 128), (2, 2048, 16, 64), (4, 512, 8, 128)]:
        for interleaved in [False, True]:
            q = torch.randn(B, S, H, D, dtype=torch.float16, device="cuda")
            k = torch.randn(B, S, H, D, dtype=torch.float16, device="cuda")
            cos = torch.randn(S, D // 2, dtype=torch.float16, device="cuda")
            sin = torch.randn(S, D // 2, dtype=torch.float16, device="cuda")
            qe_b, ke_b = baseline(q, k, cos, sin, rotary_interleaved=interleaved)
            qe_v, ke_v = v2(q, k, cos, sin, rotary_interleaved=interleaved)
            abs_diff_q = (qe_b - qe_v).abs().max().item()
            abs_diff_k = (ke_b - ke_v).abs().max().item()
            assert abs_diff_q < 1e-2, f"q diverges: {abs_diff_q}"
            assert abs_diff_k < 1e-2, f"k diverges: {abs_diff_k}"
            print(
                f"rotary_emb   B={B} S={S} H={H} D={D} "
                f"interleaved={interleaved}  abs_diff=({abs_diff_q:.2e}, {abs_diff_k:.2e})  ✓"
            )


def test_fused_moe_router_v2_correctness():
    import sys
    sys.path.insert(0, "src")
    from flag_gems_local.ops.fused_moe_router_v2 import (
        bitonic_sort_topk,
        prepare_moe_inputs_v2,
    )

    torch.manual_seed(42)
    for num_tokens, top_k in [(64, 4), (256, 8), (1024, 4)]:
        weights = torch.rand(num_tokens, top_k, device="cuda")
        ids = torch.randint(0, 8, (num_tokens, top_k), device="cuda")

        sorted_w, sorted_ids = bitonic_sort_topk(weights, ids)

        # Verify: weights should be descending
        for i in range(num_tokens):
            for j in range(top_k - 1):
                assert sorted_w[i, j] >= sorted_w[i, j + 1] - 1e-5, (
                    f"Not sorted: token {i} weight[{j}]={sorted_w[i,j]} > weight[{j+1}]={sorted_w[i,j+1]}"
                )

        # Verify: expert ids are a permutation of original
        orig_ids = torch.sort(ids, dim=1).values
        new_ids = torch.sort(sorted_ids, dim=1).values
        assert torch.equal(orig_ids, new_ids), "ids not a permutation!"

        print(f"fused_moe_router  tokens={num_tokens} top_k={top_k}  ✓")


if __name__ == "__main__":
    test_gelu_and_mul_v2_correctness()
    test_rotary_embedding_v2_correctness()
    test_fused_moe_router_v2_correctness()
    print("\nAll tests passed ✓")