# ruff: noqa
import torch
import numpy as np
from typing import Optional
from aiter.ops.tilelang.fp8_index import act_quant, fp8_index


def ref_act_quant(x: torch.Tensor, block_size: int = 128, scale_fmt: Optional[str] = None):
    """
    Reference implementation of act_quant using PyTorch.
    
    Args:
        x: Input tensor (BF16)
        block_size: Block size for quantization
        scale_fmt: Optional scale format (if not None, uses rounded scale)
                   Note: The actual value of scale_fmt doesn't matter, only whether it's None.
                   If not None, uses fast_round_scale (rounds scale to power of 2).
    
    Returns:
        quantized: FP8 quantized tensor
        scales: Scale factors (FP32)
    """
    assert x.is_contiguous(), "Input tensor must be contiguous"
    assert x.size(-1) % block_size == 0, f"Last dimension must be divisible by block_size"
    
    # Reshape to (..., M, block_size) where M = N // block_size
    shape = x.shape
    N = shape[-1]
    M = N // block_size
    x_reshaped = x.view(-1, M, block_size)
    
    # Compute absmax for each block
    # Note: TileLang kernel uses FP32 (scale_dtype) for amax computation,
    # so we should also use FP32 here to match the precision
    # Convert to FP32 before computing max to match TileLang's precision
    x_reshaped_fp32 = x_reshaped.float()  # Convert to FP32 for computation
    amax = torch.abs(x_reshaped_fp32).max(dim=-1, keepdim=True)[0]  # (..., M, 1) - use FP32
    amax = torch.clamp(amax, min=1e-4)
    
    # Compute scale (all in FP32 to match TileLang kernel)
    fp8_max = 448.0
    fp8_max_inv = 1.0 / fp8_max
    
    if scale_fmt is not None:
        # Round scale to power of 2: 2^ceil(log2(amax * fp8_max_inv))
        # This matches fast_round_scale in the kernel
        log2_scale = torch.ceil(torch.log2(amax * fp8_max_inv))
        scales = torch.pow(2.0, log2_scale)
    else:
        # Default: linear scale
        scales = amax * fp8_max_inv
    
    # scales is now FP32, matching TileLang's scale_dtype
    
    # Quantize: clamp(x / scale, -fp8_max, fp8_max)
    # Use FP32 x_reshaped_fp32 for division to match TileLang's precision
    # scales shape: (..., M, 1), x_reshaped_fp32 shape: (..., M, block_size)
    # Broadcasting should work correctly
    quantized = (x_reshaped_fp32 / scales).clamp(-fp8_max, fp8_max)
    quantized = quantized.to(torch.float8_e4m3fn)
    
    # Reshape back
    quantized = quantized.view(shape)
    scales = scales.squeeze(-1).view(*shape[:-1], M)
    
    return quantized, scales


def ref_fp8_index(
    q: torch.Tensor,
    q_s: torch.Tensor,
    k: torch.Tensor,
    k_s: torch.Tensor,
) -> torch.Tensor:
    """
    Reference implementation of fp8_index using PyTorch.
    
    Args:
        q: Query tensor (FP8), shape (b, m, h, d)
        q_s: Query weights/scale (FP32), shape (b, m, h)
        k: Key tensor (FP8), shape (b, n, d)
        k_s: Key scale (FP32), shape (b, n)
    
    Returns:
        index_score: Index scores (FP32), shape (b, m, n)
    """
    b, m, h, d = q.shape
    n = k.shape[1]
    
    # Dequantize q and k to FP32
    q_fp32 = q.float()  # (b, m, h, d)
    k_fp32 = k.float()  # (b, n, d)
    
    # Compute logits: k @ q^T
    # k: (b, n, d), q: (b, m, h, d)
    # We need to compute for each (b, m, h): k[b, n, d] @ q[b, m, h, d]^T
    # Result: (b, m, n, h)
    logits = torch.einsum('bnd,bmhd->bmnh', k_fp32, q_fp32)
    
    # Apply ReLU and multiply by q_s (weights)
    # q_s: (b, m, h) -> (b, m, 1, h) for broadcasting
    logits = torch.relu(logits)  # (b, m, n, h)
    q_s_expanded = q_s.unsqueeze(2)  # (b, m, 1, h)
    logits = logits * q_s_expanded  # (b, m, n, h)
    
    # Sum over h dimension
    logits_sum = logits.sum(dim=-1)  # (b, m, n)
    
    # Multiply by k_s
    # k_s: (b, n) -> (b, 1, n) for broadcasting
    k_s_expanded = k_s.unsqueeze(1)  # (b, 1, n)
    index_score = logits_sum * k_s_expanded  # (b, m, n)
    
    return index_score


def test_act_quant(
    B=2,
    M=4,
    N=256,
    block_size=128,
    scale_fmt=None,
    dtype=torch.bfloat16,
    check_correctness=True,
    atol=1e-3,
    rtol=1e-2,
):
    """
    Test act_quant function.
    
    Note on scale_fmt:
    - scale_fmt is only used as a boolean flag: if not None, enables rounded scale
    - The actual value (e.g., "e8m0") doesn't matter, only whether it's None
    - When scale_fmt is not None, the kernel uses fast_round_scale which rounds
      the scale to the nearest power of 2 (2^ceil(log2(scale)))
    - This can improve performance in some cases and is the "default" format when enabled
    """
    torch.random.manual_seed(0)
    device = "cuda"
    
    # Create input tensor
    x = torch.randn((B, M, N), dtype=dtype, device=device)
    
    scale_fmt_str = scale_fmt if scale_fmt else "None (linear scale)"
    print(f"Testing act_quant: B={B}, M={M}, N={N}, block_size={block_size}, scale_fmt={scale_fmt_str}")
    
    # TileLang implementation
    x_contiguous = x.contiguous()
    tl_quantized, tl_scales = act_quant(x_contiguous, block_size=block_size, scale_fmt=scale_fmt)
    
    # Reference implementation
    ref_quantized, ref_scales = ref_act_quant(x_contiguous, block_size=block_size, scale_fmt=scale_fmt)
    
    if check_correctness:
        # Check scales
        scale_atol = max(atol, 2e-4)
        scale_rtol = max(rtol, 1e-2)
        torch.testing.assert_close(
            tl_scales, ref_scales, atol=scale_atol, rtol=scale_rtol)
        print("✓ Scales match")
        
        # Check quantized values
        tl_quantized_fp32 = tl_quantized.float()
        ref_quantized_fp32 = ref_quantized.float()
        
        max_abs_diff = (tl_quantized_fp32 - ref_quantized_fp32).abs().max().item()
        max_rel_diff = ((tl_quantized_fp32 - ref_quantized_fp32).abs() / 
                       (ref_quantized_fp32.abs() + 1e-8)).max().item()
        
        torch.testing.assert_close(
            tl_quantized_fp32, ref_quantized_fp32, atol=atol, rtol=rtol)
        print(f"✓ Quantized values match (max abs diff: {max_abs_diff:.4f}, max rel diff: {max_rel_diff:.4f})")
        
        # Verify dequantization: compare TileLang's dequantized result with reference implementation's
        # Reshape for block-wise dequantization
        shape = x.shape
        M_blocks = N // block_size
        
        # TileLang dequantized result
        tl_quantized_reshaped = tl_quantized_fp32.view(-1, M_blocks, block_size)
        tl_scales_reshaped = tl_scales.view(-1, M_blocks).unsqueeze(-1)  # (..., M, 1)
        tl_dequantized = (tl_quantized_reshaped * tl_scales_reshaped).view(shape)
        
        # Reference dequantized result
        ref_quantized_reshaped = ref_quantized_fp32.view(-1, M_blocks, block_size)
        ref_scales_reshaped = ref_scales.view(-1, M_blocks).unsqueeze(-1)  # (..., M, 1)
        ref_dequantized = (ref_quantized_reshaped * ref_scales_reshaped).view(shape)
        
        # Check that dequantized values match between TileLang and reference
        torch.testing.assert_close(
            tl_dequantized, ref_dequantized, atol=atol, rtol=rtol)
        print("✓ Dequantization check passed")
    
    print("act_quant test passed!\n")


def test_fp8_index(
    b=2,
    m=4,
    n=512,
    h=16,
    d=128,
    block_size=128,
    scale_fmt=None,
    dtype=torch.bfloat16,
    check_correctness=True,
    atol=1e-2,
    rtol=1e-2,
):
    """Test fp8_index function."""
    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)
    
    # q_scale shape: (b, m, h, d//block_size)
    # k_scale shape: (b, n, d//block_size)
    
    # For fp8_index, we need:
    # q_s: (b, m, h) - weights per head
    # k_s: (b, n) - scale per position (take first block's scale)
    k_scale_index = k_scale[..., 0]  # (b, n) - take first block's scale
    
    # For weights, simulate model.py behavior:
    # weights = weights_proj(x) * n_heads**-0.5
    # weights = weights.unsqueeze(-1) * q_scale * softmax_scale
    # Since q_scale is (b, m, h, d//block_size), we need to handle this
    # For simplicity in testing, we'll use the 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)
    
    print(f"Testing fp8_index: b={b}, m={m}, n={n}, h={h}, d={d}, block_size={block_size}")
    
    # TileLang implementation
    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()
    
    tl_index_score = fp8_index(
        q_fp8_contiguous,
        q_s_contiguous,
        k_fp8_contiguous,
        k_scale_index_contiguous,
    )
    
    # Reference implementation
    ref_index_score = ref_fp8_index(
        q_fp8_contiguous,
        q_s_contiguous,
        k_fp8_contiguous,
        k_scale_index_contiguous,
    )
    
    if check_correctness:
        # Check index scores
        # Note: Due to FP8 quantization and different computation order, 
        # we need larger tolerance
        torch.testing.assert_close(
            tl_index_score, ref_index_score, atol=atol, rtol=rtol)
        print("✓ Index scores match")
    
    print("fp8_index test passed!\n")


if __name__ == "__main__":
    device_id = 0
    torch.cuda.set_device(device_id)
    
    print("=" * 80)
    print("Testing act_quant")
    print("=" * 80)
    
    # Test act_quant with different configurations
    test_act_quant(B=2, M=4, N=256, block_size=128, scale_fmt=None)
    test_act_quant(B=4, M=8, N=512, block_size=128, scale_fmt=None)
    test_act_quant(B=2, M=4, N=256, block_size=128, scale_fmt="e8m0")
    test_act_quant(B=1, M=1, N=128, block_size=128, scale_fmt=None)
    test_act_quant(B=8, M=16, N=1024, block_size=128, scale_fmt=None)
    
    print("=" * 80)
    print("Testing fp8_index")
    print("=" * 80)
    
    # Test fp8_index with different configurations
    test_fp8_index(b=2, m=4, n=256, h=16, d=128, block_size=128, scale_fmt=None)
    test_fp8_index(b=4, m=8, n=512, h=32, d=128, block_size=128, scale_fmt=None)
    test_fp8_index(b=2, m=4, n=128, h=32, d=128, block_size=128, scale_fmt=None)
    test_fp8_index(b=1, m=1, n=256, h=16, d=128, block_size=128, scale_fmt=None)
    
    print("=" * 80)
    print("All tests passed!")
    print("=" * 80)

