"""
NVFP4-1.1: BF16→FP4 quantization kernel (CuTeDSL, Blackwell SM100).

Uses cute.arch.load/store with pointer arithmetic for GMEM access.

Grid: (M, 1, 1) — 1 CTA per row, 128 threads per CTA.

Run: ~/.openclaw_workspace/fire_b200_test tests/unit/test_nvfp4_quant_kernel.py
"""
import torch
import cutlass
import cutlass.cute as cute
from cutlass import Float32, BFloat16, Float8E4M3FN, Int32, Uint8, const_expr
import cuda.bindings.driver as cuda
import cutlass.torch as ct

from dsv4.ops.quantize import quantize_activation_nvfp4, SF_VEC_SIZE


def _fmax(a, b):
    return (a + b + cute.math.absf(a - b)) / Float32(2.0)

def _fmin(a, b):
    return (a + b - cute.math.absf(a - b)) / Float32(2.0)

def _clamp(x, lo, hi):
    return _fmin(_fmax(x, lo), hi)


class Nvfp4QuantizeKernel:
    def __init__(self, block_size=16):
        self.block_size = block_size
        
    @cute.jit
    def __call__(self, x_bf16, x_fp4, x_sf, M, N, stream):
        x_bf16_ptr = x_bf16.iterator
        x_fp4_ptr = x_fp4.iterator
        x_sf_ptr = x_sf.iterator
        stride0 = x_bf16.stride[0]
        stride1 = x_bf16.stride[1]
        
        self._kernel(x_bf16_ptr, x_fp4_ptr, x_sf_ptr, M, N, stride0, stride1).launch(
            grid=(M, 1, 1), block=[32, 1, 1], stream=stream
        )
    
    @cute.kernel
    def _kernel(self, x_bf16_ptr, x_fp4_ptr, x_sf_ptr, M, N, stride0, stride1):
        tidx, _, _ = cute.arch.thread_idx()
        bidx, _, _ = cute.arch.block_idx()
        row = bidx
        
        bs = Int32(self.block_size)
        n_blocks = N // bs
        threads = Int32(32)
        blocks_per_thread = n_blocks // threads
        
        for b in cutlass.range(blocks_per_thread):
            block_idx = tidx * blocks_per_thread + b
            col_start = block_idx * bs
            
            # Pass 1: compute amax
            amax = Float32(0.0)
            for i in cutlass.range(self.block_size):
                offset = row * stride0 + (col_start + Int32(i)) * stride1
                ptr = x_bf16_ptr + offset
                raw = cute.arch.load(ptr, cutlass.Uint16)
                val = raw.bitcast(BFloat16).to(Float32)
                amax = _fmax(amax, cute.math.absf(val))
            
            scale = amax / Float32(6.0)
            # Write FP8 scale (row-major: row * n_blocks + block_idx)
            sf_offset = row * n_blocks + block_idx
            sf_val = scale.to(Float8E4M3FN).bitcast(cutlass.Uint8)
            cute.arch.store(x_sf_ptr + sf_offset, sf_val)
            
            # Pass 2: quantize and pack
            for i in cutlass.range(0, self.block_size, 2):
                off0 = row * stride0 + (col_start + Int32(i)) * stride1
                off1 = row * stride0 + (col_start + Int32(i + 1)) * stride1
                raw0 = cute.arch.load(x_bf16_ptr + off0, cutlass.Uint16)
                raw1 = cute.arch.load(x_bf16_ptr + off1, cutlass.Uint16)
                val0 = raw0.bitcast(BFloat16).to(Float32)
                val1 = raw1.bitcast(BFloat16).to(Float32)
                
                s0 = val0 / scale
                s1 = val1 / scale
                a0 = cute.math.absf(s0)
                a1 = cute.math.absf(s1)
                
                hs0 = _clamp(a0 * Float32(2.0) + Float32(0.5), Float32(0.0), Float32(12.0))
                hs1 = _clamp(a1 * Float32(2.0) + Float32(0.5), Float32(0.0), Float32(12.0))
                
                idx0 = Int32(hs0) // Int32(2)
                idx1 = Int32(hs1) // Int32(2)
                idx0 = idx0 if idx0 < Int32(7) else Int32(6)
                idx1 = idx1 if idx1 < Int32(7) else Int32(6)
                
                sign0 = Int32(1) if val0 < Float32(0.0) else Int32(0)
                sign1 = Int32(1) if val1 < Float32(0.0) else Int32(0)
                
                nibble0 = idx0 | (sign0 << Int32(3))
                nibble1 = idx1 | (sign1 << Int32(3))
                packed = nibble0 | (nibble1 << Int32(4))
                
                fp4_offset = row * (N // Int32(2)) + block_idx * Int32(self.block_size // 2) + Int32(i // 2)
                cute.arch.store(x_fp4_ptr + fp4_offset, packed.to(Uint8))


def dequantize_nvfp4(x_fp4, block_scale, global_scale, N):
    M = x_fp4.shape[0]
    block_size = SF_VEC_SIZE
    raw = x_fp4.view(torch.uint8)
    even = raw & 0x0F
    odd = (raw >> 4) & 0x0F
    indices = torch.stack([even, odd], dim=-1).reshape(M, N)
    signs = (indices >= 8).float() * -2 + 1
    mag = indices % 8
    idx_to_val = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0],
                               dtype=torch.float32, device='cuda')
    vals = idx_to_val[mag.long()]
    x_deq = signs * vals
    block_scale_exp = block_scale.repeat_interleave(block_size, dim=-1).float()
    x_deq = x_deq * block_scale_exp * global_scale
    return x_deq.to(torch.bfloat16)


def test_nvfp4_python():
    print("\n=== NVFP4 Python Round-Trip ===")
    torch.manual_seed(42)
    M, N = 128, 512
    x = torch.randn(M, N, dtype=torch.bfloat16, device='cuda')
    x_fp4, sf = quantize_activation_nvfp4(x, 1.0)
    x_deq = dequantize_nvfp4(x_fp4, sf, 1.0, N)
    cos = torch.nn.functional.cosine_similarity(
        x.flatten().float().unsqueeze(0), x_deq.flatten().float().unsqueeze(0)
    ).item()
    print(f"  Round-trip cos: {cos:.6f} ({'PASS' if cos >= 0.95 else 'FAIL'})")
    assert cos >= 0.95


def test_nvfp4_kernel():
    print("\n=== NVFP4 Kernel Quantization Test ===")
    torch.manual_seed(42)
    M, N = 128, 512
    
    x = torch.randn(M, N, dtype=torch.bfloat16, device='cuda')
    x_fp4_ref, sf_ref = quantize_activation_nvfp4(x, 1.0)
    x_deq_ref = dequantize_nvfp4(x_fp4_ref, sf_ref, 1.0, N)
    
    kernel = Nvfp4QuantizeKernel(block_size=16)
    
    x_fp4_out = torch.zeros(M, N // 2, dtype=torch.uint8, device='cuda')
    sf_out = torch.zeros(M, N // 16, dtype=torch.float8_e4m3fn, device='cuda')
    
    stream = cuda.CUstream(0)
    
    x_bf16_cute = ct.from_dlpack(x)
    x_fp4_cute = ct.from_dlpack(x_fp4_out)
    sf_cute = ct.from_dlpack(sf_out)
    
    kernel(x_bf16_cute, x_fp4_cute, sf_cute, Int32(M), Int32(N), stream)
    torch.cuda.synchronize()
    
    print(f"  FP4 nonzero: {(x_fp4_out > 0).sum().item()} / {x_fp4_out.numel()}")
    print(f"  SF nonzero: {(sf_out.float() > 0).sum().item()} / {sf_out.numel()}")
    print(f"  FP4 sample: {x_fp4_out[0, :8]}")
    print(f"  SF sample: {sf_out[0, :4].float()}")
    
    x_deq_kernel = dequantize_nvfp4(x_fp4_out, sf_out, 1.0, N)
    
    cos_kernel = torch.nn.functional.cosine_similarity(
        x.flatten().float().unsqueeze(0), x_deq_kernel.flatten().float().unsqueeze(0)
    ).item()
    cos_ref = torch.nn.functional.cosine_similarity(
        x.flatten().float().unsqueeze(0), x_deq_ref.flatten().float().unsqueeze(0)
    ).item()
    
    print(f"  Reference Python cos: {cos_ref:.6f}")
    print(f"  Kernel cos: {cos_kernel:.6f}")
    
    if cos_kernel >= 0.95:
        print(f"  ✅ Kernel quantization PASS (cos={cos_kernel:.4f})")
    else:
        print(f"  ❌ Kernel quantization FAIL (cos={cos_kernel:.4f})")
    
    fp4_match = (x_fp4_out == x_fp4_ref.view(torch.uint8)).float().mean().item()
    sf_match = (sf_out.float() == sf_ref.float()).float().mean().item()
    print(f"  FP4 byte match rate: {fp4_match:.4f}")
    print(f"  FP8 scale match rate: {sf_match:.4f}")


def test():
    print("=== NVFP4-1.1: BF4→FP4 Quantization Kernel ===")
    test_nvfp4_python()
    test_nvfp4_kernel()


if __name__ == '__main__':
    test()