"""Test NVFP4 fused router kernel against the reference path.

Phase 1: Verify reference path (BF16 linear + activation_topk) works.
Phase 2: Test CuTeDSL fused kernel (needs B200 for CuTeDSL compilation).
"""

import sys
import os
import torch

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))


def test_reference_router():
    """Test the reference BF16 linear + activation_topk path."""
    torch.manual_seed(42)
    device = "cuda"
    M = 4        # tokens
    K = 7168     # hidden size
    N = 384      # num experts
    top_k = 6
    routed_scaling_factor = 0.5

    # Create BF16 hidden states and weight
    hidden_states = torch.randn(M, K, dtype=torch.bfloat16, device=device)
    W_gate = torch.randn(K, N, dtype=torch.bfloat16, device=device)
    e_bias = torch.randn(N, dtype=torch.float32, device=device)

    # Reference: BF16 linear + activation_topk
    logits = torch.nn.functional.linear(hidden_states.float(), W_gate.T.float())
    from dsv4.kernels.router._activation_topk import run_fused_activation_topk
    out_w = torch.empty(M, top_k, dtype=torch.float32, device=device)
    out_ids = torch.empty(M, top_k, dtype=torch.int32, device=device)
    run_fused_activation_topk(logits, e_bias, routed_scaling_factor, top_k, out_w, out_ids)

    # Verify results
    print(f"Reference router test (M={M}, K={K}, N={N}, top_k={top_k}):")
    print(f"  Top-k IDs (row 0): {out_ids[0].tolist()}")
    print(f"  Top-k weights (row 0): {[f'{w:.4f}' for w in out_w[0].tolist()]}")

    # Verify: weights sum to routed_scaling_factor (approximately)
    w_sum = out_w.sum(dim=1)
    expected = routed_scaling_factor
    for row in range(M):
        diff = abs(w_sum[row].item() - expected)
        assert diff < 0.01, f"Row {row}: weight sum {w_sum[row].item():.4f} != {expected:.4f}"
    print(f"  Weight sums: {[f'{s:.4f}' for s in w_sum.tolist()]} (expected {expected})")

    # Verify: IDs are valid (0 <= id < N)
    assert (out_ids >= 0).all() and (out_ids < N).all(), "Invalid expert IDs"
    print(f"  All expert IDs in [0, {N}) ✓")

    # Verify: no duplicate IDs within each row
    for row in range(M):
        row_ids = out_ids[row].tolist()
        assert len(set(row_ids)) == len(row_ids), f"Row {row} has duplicate IDs: {row_ids}"
    print(f"  No duplicate IDs ✓")

    # Verify: weights are non-negative
    assert (out_w >= 0).all(), "Negative weights"
    print(f"  All weights non-negative ✓")

    print("Reference router test PASSED ✓")


def test_nvfp4_fused_router_import():
    """Test that the fused router kernel module imports without error."""
    try:
        from dsv4.kernels.router.nvfp4_fused_router_kernel import Nvfp4FusedRouterKernel
        print(f"Nvfp4FusedRouterKernel imported successfully")
        kernel = Nvfp4FusedRouterKernel(top_k=6)
        print(f"  mma_tiler_mn: {kernel.mma_tiler_mn}")
        print(f"  threads_per_cta: {kernel.threads_per_cta}")
        print(f"  sf_vec_size: {kernel.sf_vec_size}")
        print(f"  top_k: {kernel.top_k}")
        print("Fused router kernel class construction PASSED ✓")
    except Exception as e:
        print(f"Fused router kernel import failed: {e}")
        print("This is expected if CuTeDSL is not available.")


if __name__ == "__main__":
    test_reference_router()
    print()
    test_nvfp4_fused_router_import()