# ruff: noqa
import sys
import os
import torch
import numpy as np
from typing import Optional
from tilelang.profiler import do_bench_cudagraph, do_bench
from aiter.ops.tilelang.fp8_index import act_quant, fp8_index

# Add parent directory to path to import test functions
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../tilelang_tests'))
from test_fp8_index import ref_act_quant, ref_fp8_index


def benchmark_act_quant(
    B=2,
    M=4,
    N=256,
    block_size=128,
    scale_fmt=None,
    dtype=torch.bfloat16,
    check_correctness=False,
):
    """
    Benchmark act_quant function.
    
    Args:
        B: Batch size
        M: Sequence length
        N: Feature dimension (must be divisible by block_size)
        block_size: Block size for quantization
        scale_fmt: Optional scale format (if not None, uses rounded scale)
        dtype: Input data type
        check_correctness: Whether to check correctness against reference
    """
    torch.random.manual_seed(0)
    device = "cuda"
    
    # Create input tensor
    x = torch.randn((B, M, N), dtype=dtype, device=device)
    x_contiguous = x.contiguous()

    # Benchmark
    def fn():
        return act_quant(x_contiguous, block_size=block_size, scale_fmt=scale_fmt)
    
    ms = do_bench_cudagraph(fn)
    # ms = do_bench(fn)
    
    # Calculate performance metrics
    # Input: BF16 (2 bytes), Output: FP8 (1 byte) + FP32 scales (4 bytes per block)
    input_size_bytes = B * M * N * 2  # BF16 = 2 bytes
    output_fp8_size_bytes = B * M * N * 1  # FP8 = 1 byte
    num_blocks = N // block_size
    output_scale_size_bytes = B * M * num_blocks * 4  # FP32 = 4 bytes
    total_io_bytes = input_size_bytes + output_fp8_size_bytes + output_scale_size_bytes
    
    io_bandwidth_gbs = total_io_bytes / (ms * 1e-3) / 1e9
    
    scale_fmt_str = scale_fmt if scale_fmt else "None (linear scale)"
    print(f"{B=} {M=} {N=} {block_size=} scale_fmt={scale_fmt_str} {dtype=}")
    print(f"Average time: {ms:.3f} ms")
    print(f"IO bandwidth: {io_bandwidth_gbs:.3f} GB/s")
    print(f"  Input size: {input_size_bytes / 1e6:.2f} MB (BF16)")
    print(f"  Output size: {(output_fp8_size_bytes + output_scale_size_bytes) / 1e6:.2f} MB (FP8 + FP32 scales)")
    
    if check_correctness:
        tl_quantized, tl_scales = act_quant(x_contiguous, block_size=block_size, scale_fmt=scale_fmt)
        ref_quantized, ref_scales = ref_act_quant(x_contiguous, block_size=block_size, scale_fmt=scale_fmt)
        
        scale_atol = max(1e-3, 2e-4)
        scale_rtol = max(1e-2, 1e-2)
        torch.testing.assert_close(tl_scales, ref_scales, atol=scale_atol, rtol=scale_rtol)
        
        tl_quantized_fp32 = tl_quantized.float()
        ref_quantized_fp32 = ref_quantized.float()
        torch.testing.assert_close(tl_quantized_fp32, ref_quantized_fp32, atol=1e-3, rtol=1e-2)
        print("✓ Correctness check passed")


def benchmark_fp8_index(
    b=2,
    m=4,
    n=512,
    h=16,
    d=128,
    block_size=128,
    scale_fmt=None,
    dtype=torch.bfloat16,
    check_correctness=False,
    quiet=False,
):
    """
    Benchmark fp8_index function.
    
    Args:
        b: Batch size
        m: Query sequence length
        n: Key sequence length
        h: Number of heads
        d: Feature dimension per head (must be divisible by block_size)
        block_size: Block size for quantization
        scale_fmt: Optional scale format (if not None, uses rounded scale)
        dtype: Input data type
        check_correctness: Whether to check correctness against reference
    """
    torch.random.manual_seed(0)
    device = "cuda"
    
    # Create input tensors
    q = torch.randn((b, m, h, d), dtype=dtype, device=device)
    k = torch.randn((b, n, d), dtype=dtype, device=device)
    
    # Quantize q and k
    q_contiguous = q.contiguous()
    k_contiguous = k.contiguous()
    
    q_fp8, q_scale = act_quant(q_contiguous, block_size=block_size, scale_fmt=scale_fmt)
    k_fp8, k_scale = act_quant(k_contiguous, block_size=block_size, scale_fmt=scale_fmt)
    
    # Prepare q_s and k_s for fp8_index
    k_scale_index = k_scale[..., 0]  # (b, n) - take first block's scale
    q_scale_first = q_scale[..., 0]  # (b, m, h)
    
    # Create random weights similar to model.py
    weights_base = torch.randn((b, m, h), dtype=torch.float32, device=device) * 0.1
    softmax_scale = 1.0 / np.sqrt(h)  # Similar to n_heads**-0.5
    q_s = weights_base * q_scale_first * softmax_scale  # (b, m, h)
    
    # Make contiguous
    q_fp8_contiguous = q_fp8.contiguous()
    k_fp8_contiguous = k_fp8.contiguous()
    q_s_contiguous = q_s.contiguous()
    k_scale_index_contiguous = k_scale_index.contiguous()

    # Benchmark
    def fn():
        return fp8_index(
            q_fp8_contiguous,
            q_s_contiguous,
            k_fp8_contiguous,
            k_scale_index_contiguous,
        )
    
    ms = do_bench_cudagraph(fn)
    # ms = do_bench(fn)
    
    # Calculate performance metrics
    # Input: q (FP8), q_s (FP32), k (FP8), k_s (FP32)
    # Output: index_score (FP32)
    q_size_bytes = b * m * h * d * 1  # FP8 = 1 byte
    q_s_size_bytes = b * m * h * 4  # FP32 = 4 bytes
    k_size_bytes = b * n * d * 1  # FP8 = 1 byte
    k_s_size_bytes = b * n * 4  # FP32 = 4 bytes
    output_size_bytes = b * m * n * 4  # FP32 = 4 bytes
    
    total_io_bytes = q_size_bytes + q_s_size_bytes + k_size_bytes + k_s_size_bytes + output_size_bytes
    
    # Compute TFLOPS: GEMM operation k @ q^T
    # For each (b, m, h): k[b, n, d] @ q[b, m, h, d]^T -> (n, h)
    # Total operations: b * m * h * n * d * 2 (multiply-add)
    total_flops = b * m * h * n * d * 2
    tflops = total_flops / (ms * 1e-3) / 1e12
    
    io_bandwidth_gbs = total_io_bytes / (ms * 1e-3) / 1e9
    
    scale_fmt_str = scale_fmt if scale_fmt else "None (linear scale)"
    if not quiet:
        print(f"{b=} {m=} {n=} {h=} {d=} {block_size=} scale_fmt={scale_fmt_str} {dtype=}")
        print(f"Average time: {ms:.3f} ms")
        print(f"TFLOPS: {tflops:.3f}")
        print(f"IO bandwidth: {io_bandwidth_gbs:.3f} GB/s")
        print(f"  Input size: {(q_size_bytes + q_s_size_bytes + k_size_bytes + k_s_size_bytes) / 1e6:.2f} MB")
        print(f"  Output size: {output_size_bytes / 1e6:.2f} MB")
    
    if check_correctness:
        tl_index_score = fp8_index(
            q_fp8_contiguous,
            q_s_contiguous,
            k_fp8_contiguous,
            k_scale_index_contiguous,
        )
        ref_index_score = ref_fp8_index(
            q_fp8_contiguous,
            q_s_contiguous,
            k_fp8_contiguous,
            k_scale_index_contiguous,
        )
        torch.testing.assert_close(tl_index_score, ref_index_score, atol=1e-2, rtol=1e-2)
        if not quiet:
            print("✓ Correctness check passed")
    return {"m": m, "n": n, "ms": ms, "tflops": tflops, "gb_s": io_bandwidth_gbs}


if __name__ == "__main__":
    device_id = 0
    torch.cuda.set_device(device_id)
    
    # print("=" * 80)
    # print("Benchmarking act_quant")
    # print("=" * 80)
    
    # # Benchmark act_quant with different configurations
    benchmark_act_quant(B=1, M=16, N=512, block_size=128, scale_fmt=None, check_correctness=True)
    benchmark_act_quant(B=1, M=32, N=512, block_size=128, scale_fmt=None, check_correctness=True)
    benchmark_act_quant(B=1, M=64, N=512, block_size=128, scale_fmt=None, check_correctness=True)
    benchmark_act_quant(B=1, M=128, N=512, block_size=128, scale_fmt="e8m0", check_correctness=True)
    
    print("\n" + "=" * 80)
    print("Benchmarking fp8_index")
    print("=" * 80)

    # Original benchmark cases (commented out)
    # benchmark_fp8_index(b=1, m=1, n=2048, h=32, d=128, block_size=128, scale_fmt=None, check_correctness=True)
    # benchmark_fp8_index(b=1, m=2, n=2048, h=32, d=128, block_size=128, scale_fmt=None, check_correctness=True)
    # benchmark_fp8_index(b=1, m=4, n=2048, h=32, d=128, block_size=128, scale_fmt=None, check_correctness=True)
    # benchmark_fp8_index(b=1, m=16, n=2048, h=32, d=128, block_size=128, scale_fmt=None, check_correctness=True)
    # benchmark_fp8_index(b=1, m=32, n=2048, h=32, d=128, block_size=128, scale_fmt=None, check_correctness=True)
    # benchmark_fp8_index(b=1, m=64, n=2048, h=32, d=128, block_size=128, scale_fmt=None, check_correctness=True)
    # benchmark_fp8_index(b=1, m=128, n=2048, h=32, d=128, block_size=128, scale_fmt=None, check_correctness=True)
    # benchmark_fp8_index(b=1, m=256, n=2048, h=32, d=128, block_size=128, scale_fmt=None, check_correctness=True)
    # benchmark_fp8_index(b=1, m=512, n=2048, h=32, d=128, block_size=128, scale_fmt=None, check_correctness=True)
    # benchmark_fp8_index(b=1, m=2048, n=2048, h=32, d=128, block_size=128, scale_fmt=None, check_correctness=True)
    # benchmark_fp8_index(b=1, m=4096, n=4096, h=32, d=128, block_size=128, scale_fmt=None, check_correctness=True)
    # benchmark_fp8_index(b=1, m=4096, n=65536, h=32, d=128, block_size=128, scale_fmt=None, check_correctness=True)
    # benchmark_fp8_index(b=1, m=4088, n=65500, h=32, d=128, block_size=128, scale_fmt=None, check_correctness=True)
    # benchmark_fp8_index(b=1, m=4096, n=131072, h=32, d=128, block_size=128, scale_fmt=None, check_correctness=True)
    # benchmark_fp8_index(b=1, m=4097, n=131072, h=32, d=128, block_size=128, scale_fmt=None, check_correctness=True)
    # benchmark_fp8_index(b=1, m=4097, n=122881, h=32, d=128, block_size=128, scale_fmt=None, check_correctness=True)
    
    # Benchmark fp8_index: m in [1, 4096], n in [64, 128k], powers of 2, n >= m (no golden check)
    POW2_M = [2**i for i in range(13)]  # 1, 2, 4, ..., 4096
    POW2_N = [2**i for i in range(6, 18)]  # 64, 128, ..., 131072
    cases = [(m, n) for m in POW2_M for n in POW2_N if n >= m]
    results = []
    for idx, (m, n) in enumerate(cases):
        try:
            r = benchmark_fp8_index(
                b=1, m=m, n=n, h=32, d=128,
                block_size=128, scale_fmt=None, check_correctness=True, quiet=True
            )
            results.append(r)
            print(f"  [{idx+1}/{len(cases)}] m={m:>6} n={n:>6}  ms={r['ms']:.4f} TFLOPS={r['tflops']:.2f}")
        except Exception as e:
            print(f"  SKIP m={m} n={n}: {e}")

    # Performance statistics summary
    print("\n" + "=" * 80)
    print("fp8_index 性能统计 (b=1, h=32, d=128)")
    print("=" * 80)
    print(f"{'m':>8} {'n':>8} {'ms':>10} {'TFLOPS':>10} {'GB/s':>10}")
    print("-" * 58)
    for r in results:
        print(f"{r['m']:>8} {r['n']:>8} {r['ms']:>10.4f} {r['tflops']:>10.2f} {r['gb_s']:>10.2f}")
    print("-" * 58)
    print(f"Total: {len(results)} cases completed")
    print("=" * 80)

