"""mHC (hyper-connection) ops through the vllm-neuron#71 NKI kernels (issue #13).

    DeepInfer op (v4_modules/mhc_ops.py, hc.py)          kernel (nki/)
    expand_to_mhc                                         hc_expand.kernel
    mhc_pre_apply_mix (hc_pre collapse)                   hc_collapse.kernel
    mhc_post                                              hc_attn_out.kernel
    mhc_pre_norm_fn + split_mixes + sinkhorn (shifted)    fp32 fn GEMM (torch) + hc_mixes.kernel
    mhc_pre_big_fuse (non-shifted)                        the same, then hc_collapse.kernel

DeepInfer's semantics are kept: the ``fn`` GEMM stays fp32 (``x.float() @ fn.T``, as the ``tf32_hc_prenorm_gemm``
shim; #71 narrows ``fn`` to bf16), the rms statistic runs over the whole ``hc * hidden`` row (``hc_norm_group_size``).
Inside a compiled graph neuronx-cc runs an fp32 matmul on bf16-narrowed operands (``--auto-cast=none``
notwithstanding: the compiled shim's error equals that of ``fn.bfloat16()``, ~2e-3 of the output range), and it
folds an in-graph ``hi + lo`` split away again. So ``prepare_model`` splits every ``fn`` once, eagerly, into bf16
``hi`` + ``lo`` (``fn - hi``) and the graph runs two exact-operand matmuls with fp32 accumulation (error ~2e-6).
``hc_mixes`` bakes ``hc_sinkhorn_iters`` / ``hc_eps`` / ``rms_norm_eps`` (20 / 1e-6 / 1e-20 in V4.1), the
``2 *`` post multiplier and ``pre = sigmoid + hc_eps``; ``mixes_supported`` checks a layer's values against them
and anything else keeps the shim path. ``hc_expand`` bakes ``hc_mult = 4``.

Activations are bf16 and coefficients fp32 throughout, as in the shims. ``DEEPINFER_NEURON_NKI_HC=0`` restores the
shims.
"""

from __future__ import annotations

import os

import torch

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


def active(t: torch.Tensor, op: str | None = None) -> bool:
    return ENABLED and t.device.type == "neuron" and op not in _TORCH_OPS


#: DEEPINFER_NEURON_HC_TORCH: mHC ops that take torch even with the NKI kernels enabled. collapse (the default) is
#: bit-exact and compiles to a faster graph than its kernel. post is opt-in: post_torch is bit-equal to the kernel
#: standalone, but fused into the layer graph it rounds differently (tiny golden 5/9 on the long prompt). The mixes
#: keep the kernel path (the shim's fp32 fn GEMM runs on bf16-narrowed operands in a graph, see above).
_TORCH_OPS = set(filter(None, os.environ.get("DEEPINFER_NEURON_HC_TORCH", "collapse").split(",")))


_KERNELS: dict = {}


def _kernel(module: str):  # noqa: ANN202
    """``wrap_nki(nki.<module>.kernel)``, built once (a dict, not functools.cache, which dynamo warns about)."""
    k = _KERNELS.get(module)
    if k is None:
        from torch_neuronx.nki_hop import wrap_nki

        from deepinfer.neuron_ops.nki import hc_attn_out, hc_collapse, hc_expand, hc_mixes

        mods = {"hc_attn_out": hc_attn_out, "hc_collapse": hc_collapse, "hc_expand": hc_expand, "hc_mixes": hc_mixes}
        k = _KERNELS[module] = wrap_nki(mods[module].kernel)
    return k


def expand_supported(hidden: torch.Tensor, hc_mult: int) -> bool:
    from deepinfer.neuron_ops.nki import hc_expand

    return active(hidden) and hc_mult == hc_expand.HC


def expand(hidden: torch.Tensor, hc_mult: int) -> torch.Tensor:
    """``[..., h]`` -> ``[..., hc_mult, h]`` copies (``expand_to_mhc``)."""
    h = hidden.shape[-1]
    copies, _ = _kernel("hc_expand")(hidden.reshape(1, -1, h))
    return copies.view(*hidden.shape[:-1], hc_mult, h)


def pre_apply_mix(x: torch.Tensor, mix: torch.Tensor) -> torch.Tensor:
    """``sum_m mix[..., m] * x[..., m, :]`` in fp32, bf16 out: ``x [..., hc, h]`` bf16, ``mix [..., hc, 1]`` (or
    ``[..., hc]``) fp32 -> ``[..., h]``."""
    hc, h = x.shape[-2:]
    y = _kernel("hc_collapse")(x.reshape(1, -1, hc, h), mix.reshape(1, -1, hc).to(torch.float32))
    return y.view(*x.shape[:-2], h)


def post(x: torch.Tensor, residual: torch.Tensor, post_layer_mix: torch.Tensor, comb_res_mix: torch.Tensor) -> torch.Tensor:
    """``mhc_post``: ``x [b, n, h]``, ``residual [b, n, hc, h]``, ``post [b, n, hc, 1]``, ``comb [b, n, hc, hc]`` ->
    ``out[..., j, :] = post[j] * x + sum_i comb[i, j] * residual[i]`` ``[b, n, hc, h]``."""
    b, n, hc, h = residual.shape
    return _kernel("hc_attn_out")(
        x.reshape(b, n, h), residual, post_layer_mix.reshape(b, n, hc), comb_res_mix.reshape(b, n, hc, hc)
    )


def post_torch(x: torch.Tensor, residual: torch.Tensor, post_layer_mix: torch.Tensor, comb_res_mix: torch.Tensor) -> torch.Tensor:
    """``post`` in torch, in ``hc_attn_out``'s fp32 operation order (``comb`` terms first, ``post * x`` last), so it
    compiles to plain elementwise work instead of the kernel (DEEPINFER_NEURON_HC_TORCH=post)."""
    b, n, hc, h = residual.shape
    res = residual.float()
    xf = x.reshape(b, n, 1, h).float()
    comb = comb_res_mix.reshape(b, n, hc, hc)
    acc = res[:, :, 0:1, :] * comb[:, :, 0, :].unsqueeze(-1)  # [b, n, j, h]: comb[0, j] * res[0]
    for i in range(1, hc):
        acc = res[:, :, i : i + 1, :] * comb[:, :, i, :].unsqueeze(-1) + acc
    out = xf * post_layer_mix.reshape(b, n, hc, 1) + acc
    return out.to(residual.dtype)


def split_fn(fn: torch.Tensor) -> None:
    """Attach ``fn``'s bf16 ``hi`` / ``lo`` halves (``hi + lo == fn`` to ~2**-17); eager only, see module doc."""
    hi = fn.detach().to(torch.bfloat16)
    fn._nki_fn_hi = hi
    fn._nki_fn_lo = (fn.detach() - hi.float()).to(torch.bfloat16)


def prepare_model(model: torch.nn.Module) -> int:
    """``split_fn`` every mHC ``fn`` of the model (``hc_attn_fn`` / ``hc_ffn_fn``). Returns the count."""
    if not ENABLED:
        return 0
    count = 0
    for name, p in model.named_parameters():
        if name.endswith(("hc_attn_fn", "hc_ffn_fn")) and p.dtype == torch.float32 and p.dim() == 2:
            split_fn(p)
            count += 1
    return count


def _fn_gemm(flat: torch.Tensor, fn: torch.Tensor) -> torch.Tensor:
    """``flat.float() @ fn.T`` at fp32 accuracy, also inside a compiled graph."""
    if getattr(fn, "_nki_fn_hi", None) is None and not torch.compiler.is_compiling():
        split_fn(fn)
    hi = getattr(fn, "_nki_fn_hi", None)
    xf = flat.float()
    if hi is None:  # not prepared and already tracing: bf16-narrowed fn in a graph, like the shim
        return xf @ fn.transpose(0, 1)
    return xf @ hi.float().transpose(0, 1) + xf @ fn._nki_fn_lo.float().transpose(0, 1)


def mixes_supported(
    residual: torch.Tensor,
    norm_eps: float,
    pre_eps: float,
    sinkhorn_eps: float,
    post_mult_value: float,
    sinkhorn_repeat: int,
) -> bool:
    from deepinfer.neuron_ops.nki import hc_mixes

    return (
        active(residual)
        and residual.dtype == torch.bfloat16
        and sinkhorn_repeat == hc_mixes.SINKHORN_ITERS
        and pre_eps == hc_mixes.HC_EPS
        and sinkhorn_eps == hc_mixes.HC_EPS
        and norm_eps == hc_mixes.NORM_EPS
        and post_mult_value == 2.0
    )


def mixes(
    residual: torch.Tensor, fn: torch.Tensor, scale: torch.Tensor, base: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """The layer's three mixes from the trunk ``[..., hc, h]`` bf16: ``pre [..., hc, 1]``, ``post [..., hc, 1]``
    and the Sinkhorn-normalized ``comb [..., hc, hc]``, all fp32 (``mhc_pre_norm_fn`` + ``mhc_pre_split_mixes`` +
    ``sinkhorn_normalize``). ``fn [(2 + hc) * hc, hc * h]`` fp32 already carries the norm weight."""
    hc, h = residual.shape[-2:]
    lead = residual.shape[:-2]
    flat = residual.reshape(-1, hc * h)
    raw = _fn_gemm(flat, fn)  # fp32, as tf32_hc_prenorm_gemm's shim
    pre, post_mix, comb = _kernel("hc_mixes")(raw, flat, base.reshape(1, -1), scale.reshape(1, 3))
    return pre.view(*lead, hc, 1), post_mix.view(*lead, hc, 1), comb.view(*lead, hc, hc)


def pre_big_fuse(
    residual: torch.Tensor, fn: torch.Tensor, scale: torch.Tensor, base: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """``mhc_pre_big_fuse``: ``(post [..., hc, 1], comb [..., hc, hc], layer_input [..., h])``."""
    pre, post_mix, comb = mixes(residual, fn, scale, base)
    if "collapse" in _TORCH_OPS:
        from deepinfer.model_executor.models.v4_modules.mhc_ops import mhc_pre_apply_mix

        return post_mix, comb, mhc_pre_apply_mix(residual, pre)
    return post_mix, comb, pre_apply_mix(residual, pre)