"""Pure-torch drop-in for the official DeepSeek-V4.1 reference `inference/kernel.py` (TileLang GPU kernels).

Mirrors the public API of kernel.py (act_quant, fp4_act_quant, fp8_gemm, fp4_gemm, sparse_attn,
hc_split_sinkhorn) with the same numerics, so the unmodified reference `model.py` can run on CPU:

    import ref_kernels_torch, sys
    sys.modules["kernel"] = ref_kernels_torch   # before `import model`

Numerics reproduced from kernel.py (line numbers refer to the downloaded kernel.py):
  * fast_round_scale (22-37): scale = 2**ceil(log2(amax * fp32(1/max))) taken from the fp32 bit pattern.
  * act_quant (41-124): per (row, block) amax clamped to >= 1e-4; e4m3 RNE cast after clamp(+-448);
    inplace=True writes back dtype(fp32(e4m3(x/s)) * s) (a quant->dequant round trip).
  * fp4_act_quant (128-204): e8m0 scales: amax >= 6*2**-126, pow2-ceil(amax/6); e4m3 scales: amax >= 6*2**-9,
    s = fp32(e4m3(amax / 6)); values RNE to the e2m1 grid after clamp(+-6); inplace => dtype(fp32(e2m1)*s).
  * fp8_gemm / fp4_gemm (208-307, 478-591): per 32-wide K block, fp32 partial GEMM, then
    acc += partial * scale_a[m, k] * scale_b[n(/32), k] in K-block order; output in torch default dtype.
    FP4 weights: float4_e2m1fn_x2 [N, K/2], element 2j = low nibble, 2j+1 = high nibble.
  * sparse_attn (311-403): FlashAttention-style online softmax over blocks of 64 gathered indices
    (-1 = masked, gathered as zeros with score -inf), running max initialised to -1e30, the probabilities
    are cast to bf16 before the P@V GEMM (acc_s_cast), sink added as exp(sink - max) to the denominator.
  * hc_split_sinkhorn (407-474): sigmoid pre (+eps), 2*sigmoid post, softmax(+eps) then column
    normalisation, then (iters-1) x (row, column) normalisation, each with +eps in the denominator.
"""

from __future__ import annotations

from typing import Optional

import torch

FP8_MAX = 448.0
FP4_MAX = 6.0
FP4_TABLE = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0])
_GEMM_BLOCK_K = 32  # kernel.py: block_K = group_size (32 for this model) / fp4 sub_K = 32
_ATTN_BLOCK = 64  # kernel.py sparse_attn_kernel: block = 64


def _fast_round_scale(amax: torch.Tensor, max_inv: float) -> torch.Tensor:
    """kernel.py fast_round_scale: 2**ceil(log2(amax * max_inv)), via the IEEE-754 bits."""
    v = (amax.float() * torch.tensor(max_inv, dtype=torch.float32)).contiguous()
    bits = v.view(torch.int32)
    exp = (bits >> 23) & 0xFF
    man = bits & ((1 << 23) - 1)
    e = exp - 127 + (man != 0).to(torch.int32)
    # fast_pow2: (e + 127) << 23 reinterpreted as fp32
    return ((e + 127) << 23).to(torch.int32).view(torch.float32)


def _e2m1_round(v: torch.Tensor) -> torch.Tensor:
    """fp32 in [-6, 6] -> nearest e2m1 value (round half to even), returned as fp32."""
    a = v.abs()
    code = torch.zeros_like(a, dtype=torch.int64)
    for b in (0.25, 1.25, 2.5, 5.0):  # ties go down (to the even code)
        code += a > b
    for b in (0.75, 1.75, 3.5):  # ties go up (to the even code)
        code += a >= b
    mag = FP4_TABLE.to(v.device)[code]
    return torch.where(v < 0, -mag, mag)


def _e2m1_codes(v: torch.Tensor) -> torch.Tensor:
    a = v.abs()
    code = torch.zeros_like(a, dtype=torch.int32)
    for b in (0.25, 1.25, 2.5, 5.0):
        code += (a > b).to(torch.int32)
    for b in (0.75, 1.75, 3.5):
        code += (a >= b).to(torch.int32)
    return code | ((v < 0).to(torch.int32) << 3)


def _unpack_fp4(b: torch.Tensor) -> torch.Tensor:
    """float4_e2m1fn_x2 / int8 [..., K/2] -> fp32 [..., K]; low nibble first."""
    x = b.contiguous().view(torch.uint8)
    table = FP4_TABLE.to(x.device)
    lo, hi = (x & 0x0F).long(), (x >> 4).long()
    return torch.stack([table[lo], table[hi]], dim=-1).flatten(-2)


def _to_scale_dtype(s: torch.Tensor, scale_dtype: torch.dtype) -> torch.Tensor:
    if scale_dtype == torch.float8_e8m0fnu:
        # s is a power of two here: its exponent byte is the e8m0 encoding
        return ((s.contiguous().view(torch.int32) >> 23) & 0xFF).to(torch.uint8).view(torch.float8_e8m0fnu)
    return s.to(scale_dtype)


def act_quant(
    x: torch.Tensor,
    block_size: int = 128,
    scale_fmt: Optional[str] = None,
    scale_dtype: torch.dtype = torch.float32,
    inplace: bool = False,
):
    """Block-wise FP8 quantization; inplace=True does the fused quant+dequant back to x.dtype."""
    N = x.size(-1)
    assert N % block_size == 0
    z = x.contiguous()
    xf = z.float().unflatten(-1, (N // block_size, block_size))
    amax = xf.abs().amax(-1).clamp_min(1e-4)
    if scale_fmt is not None:
        s = _fast_round_scale(amax, 1.0 / FP8_MAX)
    else:
        s = amax * torch.tensor(1.0 / FP8_MAX, dtype=torch.float32)
    q = (xf / s.unsqueeze(-1)).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
    if inplace:
        y = (q.float() * s.unsqueeze(-1)).to(x.dtype).flatten(-2)
        x.copy_(y)
        return x
    return q.flatten(-2), _to_scale_dtype(s, scale_dtype)


def fp4_act_quant(
    x: torch.Tensor,
    block_size: int = 32,
    inplace: bool = False,
    scale_dtype: torch.dtype = torch.float8_e8m0fnu,
):
    """FP4 with E8M0 scales (indexer) or E4M3 scales (compressed KV); inplace writes the dequantized values."""
    assert scale_dtype in (torch.float8_e8m0fnu, torch.float8_e4m3fn)
    N = x.size(-1)
    assert N % block_size == 0
    z = x.contiguous()
    xf = z.float().unflatten(-1, (N // block_size, block_size))
    amax = xf.abs().amax(-1)
    if scale_dtype == torch.float8_e4m3fn:
        amax = amax.clamp_min(FP4_MAX * 2.0**-9)
        s = (amax / FP4_MAX).to(torch.float8_e4m3fn).float()
    else:
        amax = amax.clamp_min(FP4_MAX * 2.0**-126)
        s = _fast_round_scale(amax, 1.0 / FP4_MAX)
    v = (xf / s.unsqueeze(-1)).clamp(-FP4_MAX, FP4_MAX)
    if inplace:
        y = (_e2m1_round(v) * s.unsqueeze(-1)).to(x.dtype).flatten(-2)
        x.copy_(y)
        return x
    codes = _e2m1_codes(v).flatten(-2)
    packed = (codes[..., 0::2] | (codes[..., 1::2] << 4)).to(torch.uint8).view(torch.float4_e2m1fn_x2)
    s_out = _to_scale_dtype(s, scale_dtype) if scale_dtype == torch.float8_e8m0fnu else s.to(torch.float8_e4m3fn)
    return packed, s_out


def _blocked_gemm(A: torch.Tensor, sa: torch.Tensor, B: torch.Tensor, sb_rows: torch.Tensor, act_group: int) -> torch.Tensor:
    """A [M,K] fp32 (exact fp8 values), sa [M, K/act_group]; B [N,K] fp32 (exact), sb_rows [N, K/32].
    acc += (A_k @ B_k^T) * sa[:, k // n_sub] * sb[:, k] for each 32-wide K block, in order."""
    M, K = A.shape
    N = B.shape[0]
    n_sub = act_group // _GEMM_BLOCK_K
    acc = torch.zeros(M, N, dtype=torch.float32, device=A.device)
    for k in range(K // _GEMM_BLOCK_K):
        sl = slice(k * _GEMM_BLOCK_K, (k + 1) * _GEMM_BLOCK_K)
        part = A[:, sl] @ B[:, sl].T
        acc += part * sa[:, k // n_sub].unsqueeze(1) * sb_rows[:, k].unsqueeze(0)
    return acc


def fp8_gemm(
    a: torch.Tensor,
    a_s: torch.Tensor,
    b: torch.Tensor,
    b_s: torch.Tensor,
    scale_dtype: torch.dtype = torch.float32,
    block_size: int = 128,
) -> torch.Tensor:
    """C[M,N] = A[M,K] @ B[N,K]^T with per-block FP8 scaling (weight scales per (N block, K block))."""
    assert a.is_contiguous() and b.is_contiguous()
    assert a_s.is_contiguous() and b_s.is_contiguous()
    assert block_size == _GEMM_BLOCK_K, "this model only uses 32x32 blocks"
    K = a.size(-1)
    M = a.numel() // K
    N = b.size(0)
    assert K % block_size == 0
    assert a_s.numel() == M * (K // block_size)
    assert b_s.shape == ((N + block_size - 1) // block_size, K // block_size)
    sb_rows = b_s.float().repeat_interleave(block_size, dim=0)[:N]
    acc = _blocked_gemm(a.view(M, K).float(), a_s.view(M, -1).float(), b.float(), sb_rows, block_size)
    return acc.to(torch.get_default_dtype()).view(*a.size()[:-1], N)


def fp4_gemm(
    a: torch.Tensor,
    a_s: torch.Tensor,
    b: torch.Tensor,
    b_s: torch.Tensor,
    scale_dtype: torch.dtype = torch.float32,
    act_block_size: int = 128,
) -> torch.Tensor:
    """C[M,N] = A_fp8[M,K] @ B_fp4[N,K]^T; B is [N, K/2] float4_e2m1fn_x2 with per-32 E8M0 scales."""
    assert a.is_contiguous() and b.is_contiguous()
    assert a_s.is_contiguous() and b_s.is_contiguous()
    K = a.size(-1)
    M = a.numel() // K
    N = b.size(0)
    assert act_block_size in (32, 128)
    assert K % act_block_size == 0
    assert a_s.numel() == M * (K // act_block_size)
    assert b_s.shape == (N, K // 32)
    acc = _blocked_gemm(a.view(M, K).float(), a_s.view(M, -1).float(), _unpack_fp4(b), b_s.float(), act_block_size)
    return acc.to(torch.get_default_dtype()).view(*a.size()[:-1], N)


def sparse_attn(q: torch.Tensor, kv: torch.Tensor, attn_sink: torch.Tensor, topk_idxs: torch.Tensor, softmax_scale: float) -> torch.Tensor:
    """q [b,m,h,d] bf16, kv [b,n,d] bf16, attn_sink [h] fp32, topk_idxs [b,m,topk] int32 (-1 = none)."""
    b, m, h, d = q.shape
    topk = topk_idxs.size(-1)
    qf = q.float()
    sink = attn_sink.float()
    scores_max = torch.full((b, m, h), -1e30, dtype=torch.float32, device=q.device)
    sum_exp = torch.zeros((b, m, h), dtype=torch.float32, device=q.device)
    acc_o = torch.zeros((b, m, h, d), dtype=torch.float32, device=q.device)
    batch_ix = torch.arange(b, device=q.device).view(b, 1, 1)
    n_blocks = (topk + _ATTN_BLOCK - 1) // _ATTN_BLOCK
    for t in range(n_blocks):
        idx = topk_idxs[..., t * _ATTN_BLOCK : (t + 1) * _ATTN_BLOCK].long()
        if idx.size(-1) < _ATTN_BLOCK:  # positions past topk read as -1
            idx = torch.cat([idx, idx.new_full((b, m, _ATTN_BLOCK - idx.size(-1)), -1)], dim=-1)
        valid = idx != -1
        kv_blk = kv[batch_ix, idx.clamp_min(0)]  # [b, m, 64, d]
        kv_blk = torch.where(valid.unsqueeze(-1), kv_blk, torch.zeros((), dtype=kv.dtype)).float()
        s = torch.einsum("bmhd,bmkd->bmhk", qf, kv_blk)
        s = torch.where(valid.unsqueeze(2), s, torch.tensor(-torch.inf))
        s = s * softmax_scale
        prev = scores_max
        scores_max = torch.maximum(prev, s.amax(-1))
        scale = torch.exp(prev - scores_max)
        p = torch.exp(s - scores_max.unsqueeze(-1))
        sum_exp = sum_exp * scale + p.sum(-1)
        acc_o = acc_o * scale.unsqueeze(-1) + torch.einsum("bmhk,bmkd->bmhd", p.to(torch.bfloat16).float(), kv_blk)
    sum_exp = sum_exp + torch.exp(sink.view(1, 1, h) - scores_max)
    return (acc_o / sum_exp.unsqueeze(-1)).to(q.dtype)


def hc_split_sinkhorn(
    mixes: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor, hc_mult: int = 4, sinkhorn_iters: int = 20, eps: float = 1e-6
):
    hc = hc_mult
    m = mixes.float()
    sc, base = hc_scale.float(), hc_base.float()
    pre = torch.sigmoid(m[..., :hc] * sc[0] + base[:hc]) + eps
    post = 2 * torch.sigmoid(m[..., hc : 2 * hc] * sc[1] + base[hc : 2 * hc])
    comb = m[..., 2 * hc :].unflatten(-1, (hc, hc)) * sc[2] + base[2 * hc :].view(hc, hc)
    comb = torch.exp(comb - comb.amax(-1, keepdim=True))
    comb = comb / comb.sum(-1, keepdim=True) + eps
    comb = comb / (comb.sum(-2, keepdim=True) + eps)
    for _ in range(sinkhorn_iters - 1):
        comb = comb / (comb.sum(-1, keepdim=True) + eps)
        comb = comb / (comb.sum(-2, keepdim=True) + eps)
    return pre, post, comb