"""Block-quantized GEMMs through the vllm-neuron#71 NKI kernels (issue #10, #25).

DeepInfer keeps fp8 weights as ``[n, k]`` e4m3 with per-row, per-32-column packed ue8m0 scales (expanded from the
checkpoint's 32x32 blocks) and fp4 experts as ``[e, n, k/2]`` e2m1 pairs with per-row scales. The kernels take the
checkpoint form: e4m3 bytes with a ``[ceil(n/32), ceil(k/32)]`` e8m0 scale, e2m1 bytes with ``[n, k/32]``. The
kernel scales are prepared once after loading (``prepare_model``) and kept on the weight as ``_nki_scale``.

Activations arrive fp8-quantized (1x32 blocks); they are dequantized to bf16 exactly (an e4m3 value times a power of
two) and fed to the kernel. ``fp4_linear`` re-quantizes them through the same 1x32 fp8 round trip, which is
idempotent on already-quantized values.
"""

from __future__ import annotations

import functools
import os

import torch

ENABLED = os.environ.get("DEEPINFER_NEURON_NKI_GEMM", "1") == "1"


def _bytes(sf: torch.Tensor) -> torch.Tensor:
    """Packed ue8m0 scales (int16 pairs, uint8 or float8_e8m0fnu) -> uint8 exponent bytes, row-major."""
    return sf.contiguous().view(torch.uint8)


def _kernel_scale(weight: torch.Tensor, sf: torch.Tensor) -> torch.Tensor | None:
    """The kernel's scale for one weight, or None if its scales do not fit the kernel's layout."""
    k_blocks = -(-(weight.shape[-1] * (2 if weight.dtype == torch.int8 else 1)) // 32)
    b = _bytes(sf)[..., :k_blocks]
    if weight.dtype == torch.int8:  # fp4 experts: per-row scales, as the kernel takes them
        return b.contiguous()
    n = weight.shape[-2]
    blocks = b[..., ::32, :]
    if not torch.equal(blocks.repeat_interleave(32, dim=-2)[..., :n, :], b):  # rows of a block must share a scale
        return None
    return blocks.contiguous()


def prepare_model(model: torch.nn.Module) -> int:
    """Attach ``_nki_scale`` to every fp8 / fp4 weight the kernels can take. Returns the count."""
    if not ENABLED:
        return 0
    count = 0
    skipped = []
    for name, p in model.named_parameters():
        sf = getattr(p, "sf", None)
        if sf is None or p.dtype not in (torch.float8_e4m3fn, torch.int8) or p.dim() not in (2, 3):
            continue
        # On the host: the scale tensors are small, and eager device ops would compile once per weight shape.
        sf_host = (sf.data if isinstance(sf, torch.nn.Parameter) else sf).cpu()
        scale = _kernel_scale(p.data, sf_host)
        if scale is not None:
            scale = scale.to(p.device)
            p._nki_scale = scale
            # The weight bytes as their own uint8 tensor (same storage): a dtype bitcast of a graph input feeding the
            # kernel crashes neuronx-cc (InferNonlocalTensors), a uint8 input does not.
            p._nki_weight = p.data.view(torch.uint8)
            count += 1
        else:
            skipped.append(f"{name}{tuple(p.shape)}")
    if skipped:
        # These stay on the deep_gemm shims: each call dequantizes the whole weight inside the graph.
        from deepinfer.logger import init_logger

        init_logger(__name__).warning(f"Neuron: {len(skipped)} block-quant weights without the NKI kernels: {skipped[:12]}")
    return count


#: At LNC=2 a logical NeuronCore is two physical cores, and a kernel launched at grid 1 has code for only one of them:
#: half the machine idles. The ``_2core`` kernels divide the output columns by ``nl.program_id(0)`` and run at grid 2
#: (vllm-neuron#71): 1.7-1.9x on every decode GEMM, bitwise equal (issue #27). DEEPINFER_NEURON_BQ_2CORE=0 turns it off.
TWO_CORE = os.environ.get("DEEPINFER_NEURON_BQ_2CORE", "1") == "1"


@functools.cache
def _kernel(name: str, grid: int = 0):  # noqa: ANN202
    from torch_neuronx.nki_hop import wrap_nki

    from deepinfer.neuron_ops.nki import blockquant

    kernel = wrap_nki(getattr(blockquant, name))
    return kernel[grid] if grid else kernel


def _linear_kernel(name: str, n: int):  # noqa: ANN202
    """``name`` at grid 2 when its output columns split into two whole 128-column halves, else at grid 1."""
    if TWO_CORE and n % 256 == 0:
        return _kernel(f"{name}_2core", 2)
    return _kernel(name)


def _activation_bf16(a: torch.Tensor, sfa: torch.Tensor, k: int) -> torch.Tensor:
    from deep_gemm._quant import decode_data, sf_to_scales

    return (decode_data(a) * sf_to_scales(sfa, k)).to(torch.bfloat16)


def supported(a: torch.Tensor, b: torch.Tensor, k: int, n: int) -> bool:
    return (
        ENABLED
        and a.device.type == "neuron"
        and getattr(b, "_nki_scale", None) is not None
        and k % 128 == 0
        and n % 32 == 0
    )


def gemm(a: torch.Tensor, sfa: torch.Tensor, b: torch.Tensor, d: torch.Tensor, group: int | None = None) -> None:
    """``d = dequant(a) @ dequant(b[group])^T`` through the NKI kernels."""
    m = a.shape[0]
    x = _activation_bf16(a, sfa, a.shape[-1])
    scale = b._nki_scale if group is None else b._nki_scale[group]
    w = b._nki_weight if group is None else b._nki_weight[group]
    name = "fp4_linear" if b.dtype == torch.int8 else "fp8_linear_raw"
    out = _linear_kernel(name, w.shape[0])(x, w, scale)
    d.copy_(out[:m].to(d.dtype))


def linear_qdq(x: torch.Tensor, w: torch.Tensor, group: int | None = None) -> torch.Tensor:
    """bf16 ``x [m, k]`` against a prepared fp8 / fp4 weight, the kernel quantizing ``x`` to fp8 (1x32, power-of-two
    scales, max 224) on the way in -- the same quantization the shims apply, so callers can skip the fp8 activation
    and its int16 scales entirely. Returns bf16 ``[m, n]``."""
    scale = w._nki_scale if group is None else w._nki_scale[group]
    weight = w._nki_weight if group is None else w._nki_weight[group]
    name = "fp4_linear" if w.dtype == torch.int8 else "fp8_linear"
    return _linear_kernel(name, weight.shape[0])(x, weight, scale)[: x.shape[0]]


def linear_qdq_raw(x: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    """``linear_qdq`` for a bare fp8 weight (uint8 bytes ``[n, k]`` + kernel scale), e.g. a rank's shared-expert
    slice (static_moe.prepare_shared_slices)."""
    return _linear_kernel("fp8_linear", weight.shape[0])(x, weight, scale)[: x.shape[0]]


def linear_qdq_gated(x: torch.Tensor, w: torch.Tensor, group: int, flag: torch.Tensor) -> torch.Tensor | None:
    """``linear_qdq`` for fp4 expert ``group`` that the kernel skips (zeros) unless ``flag`` (int32 ``[1, 1]``) is 1.
    None when the gated kernel does not apply (not fp4, or columns that do not split across the two cores)."""
    scale = w._nki_scale[group]
    weight = w._nki_weight[group]
    if w.dtype != torch.int8 or not TWO_CORE or weight.shape[0] % 256 != 0:
        return None
    return _kernel("fp4_linear_2core_gated", 2)(x, weight, scale, flag)[: x.shape[0]]


def has_kernel(w: torch.Tensor | None) -> bool:
    return ENABLED and w is not None and getattr(w, "_nki_scale", None) is not None and w.device.type == "neuron"


def bf16_norms(model: torch.nn.Module) -> int:
    """Stop the attention / ffn norms casting to fp8 where every consumer is an NKI kernel (prepare_model first).

    The kernels quantize the bf16 hidden themselves, so the norm's fp8 copy and its scales were dead -- but they
    crossed the graph boundary after the mHC pre block as outputs, and their encode (floor / select / power over the
    whole hidden, mostly on Vector and GpSimd) was about half of that block (3.6 ms of a 512-row prefill chunk,
    twice a layer). Returns the count of norms switched to bf16.
    """
    count = 0
    for layer in model.modules():
        attn, mlp = getattr(layer, "self_attn", None), getattr(layer, "mlp", None)
        pairs = (
            (getattr(layer, "attn_layernorm", None), has_kernel(getattr(getattr(attn, "q_a_kv_swa_proj", None), "weight", None))),
            (getattr(layer, "ffn_layernorm", None), has_kernel(getattr(mlp, "gate_up_proj_weight", None)) and not getattr(mlp, "use_mega_moe", False)),
        )
        for norm, nki_only in pairs:
            if nki_only and getattr(norm, "use_fp8", False):
                norm.use_fp8 = False
                count += 1
    return count