"""Reference quant/dequant ported from vendor/FlashMLA/tests/quant.py (CPU, no kernelkit)."""

import torch

META = {
    "v41_fp8": (512, 448, 64, 32, 16),
    "v41_fp4": (512, 448, 64, 16, 32),
}
BPT = {"v41_fp8": 528, "v41_fp4": 288}

_E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32)


def _cast_scale_inv_to_ue8m0(scales_inv, out_dtype=torch.float32):
    return torch.pow(2, torch.clamp_min(scales_inv, 1e-4).log2().ceil()).to(out_dtype)


def _quantize_to_e2m1(x):
    x = x.float()
    mags = _E2M1
    sign = (torch.signbit(x)).to(torch.uint8) << 3
    a = torch.nan_to_num(x.abs(), nan=0.0, posinf=6.0).clamp_max(6.0)
    mids = (mags[:-1] + mags[1:]) / 2
    code = torch.bucketize(a, mids, right=True)
    on_tie = (a.unsqueeze(-1) == mids).any(dim=-1)
    tie_code = torch.bucketize(a, mids, right=False)
    code = torch.where(on_tie, tie_code + (tie_code & 1), code)
    return (sign | code.to(torch.uint8)).to(torch.uint8)


def quantize_k_cache(input_k_cache, fmt):
    """input_k_cache: (num_blocks, block_size, 1, 512) bf16 -> uint8 (nb, bs, 1, bpt)."""
    d, d_nope, d_rope, tile_size, num_tiles = META[fmt]
    num_blocks, block_size, h_k, _ = input_k_cache.shape
    assert h_k == 1
    x = input_k_cache.squeeze(2)
    bpt = BPT[fmt]
    result = torch.empty((num_blocks, block_size, bpt), dtype=torch.uint8)

    if fmt == "v41_fp8":
        data = result[..., : d_nope + d_rope].view(torch.float8_e4m3fn)
        sf = result[..., d_nope + d_rope :]
        for t in range(16):
            sl = x[..., t * tile_size : (t + 1) * tile_size]
            inv = torch.abs(sl).max(dim=-1).values.float() / 448.0
            inv = _cast_scale_inv_to_ue8m0(inv)
            # ue8m0 byte = bf16 bits >> 7 (Ascend path of tests/quant.py)
            sf[:, :, t] = (inv.view(torch.int32) // 2**23).to(torch.uint8)
            data[:, :, t * tile_size : (t + 1) * tile_size] = (sl.float() / inv.unsqueeze(-1)).to(
                torch.float8_e4m3fn
            )
    else:
        data = result[..., : d // 2]
        sf = result[..., d // 2 :].view(torch.float8_e4m3fn)
        xf = x.float()
        amax = torch.nan_to_num(xf.abs(), nan=float("inf")).view(num_blocks, block_size, num_tiles, tile_size).amax(-1)
        scale = torch.clamp(amax / 6.0, 2.0**-9, 448.0).to(torch.float8_e4m3fn)
        normalized = xf.view(num_blocks, block_size, num_tiles, tile_size) / scale.float().unsqueeze(-1)
        codes = _quantize_to_e2m1(normalized).view(num_blocks, block_size, d)
        data[...] = codes[..., 0::2] | (codes[..., 1::2] << 4)
        sf[...] = scale
    return result.view(num_blocks, block_size, 1, bpt)


def dequantize_k_cache(quant_k_cache, fmt):
    d, d_nope, d_rope, tile_size, num_tiles = META[fmt]
    num_blocks, block_size, h_k, _ = quant_k_cache.shape
    assert h_k == 1
    out = torch.empty((num_blocks, block_size, d), dtype=torch.bfloat16)
    q = quant_k_cache.view(num_blocks, block_size, -1)
    if fmt == "v41_fp8":
        raw = q[..., : d_nope + d_rope].view(torch.float8_e4m3fn)
        sf = q[..., d_nope + d_rope :]
        for t in range(16):
            cur = raw[..., t * tile_size : (t + 1) * tile_size].to(torch.bfloat16)
            scales = (sf[:, :, t].to(torch.int32) * 2**7).to(torch.int16).view(torch.bfloat16)
            out[..., t * tile_size : (t + 1) * tile_size] = cur * scales.unsqueeze(-1)
    else:
        data = q[..., : d // 2]
        sf = q[..., d // 2 :].view(torch.float8_e4m3fn)
        codes = torch.empty((num_blocks, block_size, d), dtype=torch.uint8)
        codes[..., 0::2] = data & 0xF
        codes[..., 1::2] = data >> 4
        val = _E2M1[(codes & 7).long()]
        val = torch.where((codes & 8) != 0, -val, val).view(num_blocks, block_size, num_tiles, tile_size)
        out[...] = (val.to(torch.bfloat16) * sf.to(torch.bfloat16).unsqueeze(-1)).view(
            num_blocks, block_size, d
        )
    return out.view(num_blocks, block_size, 1, d)