"""MHC modeling helpers that only DeepInfer needs.

`mhc_head` and the fused inference path `mhc_pre_big_fuse` combine the MHC operators
from `mhc_ops`; both compositions are inference-only and have no other caller.
"""

from __future__ import annotations

import torch
import torch.nn.functional as F  # noqa: N812
from tile_kernels.mhc.pre_big_fuse_kernel import mhc_pre_big_fuse_fwd

from deepinfer.model_executor.models.v4_modules.mhc_ops import (
    default_n_splits,
    mhc_head_compute_mix,
    mhc_pre_apply_mix,
    mhc_pre_norm_fn,
    mhc_prenorm_gemm,
)


def mhc_head(
    residual: torch.Tensor,
    fn: torch.Tensor,
    scale: torch.Tensor,
    base: torch.Tensor,
    *,
    norm_weight: torch.Tensor | None = None,
    norm_eps: float = 1e-6,
    mhc_mult: int = 4,
    pre_eps: float = 1e-6,
    n_splits: int | None = None,
) -> torch.Tensor:
    """MHC head processing for the final language model head.

    Combines pre_norm_fn + head_compute_mix + pre_apply_mix into a single call.
    The fn parameter follows the same convention as block-level fn:
    [mhc_mult, mhc_mult * hidden_size], which is padded to [mhc_mult * (mhc_mult + 2), ...]
    internally to reuse the pre_norm_fn kernel.

    Args:
        residual: MHC residual tensor of shape [..., mhc_mult, hidden_size]
        fn: weight matrix of shape [mhc_mult, mhc_mult * hidden_size]
        scale: sigmoid scaling of shape [1] or scalar
        base: mix biases of shape [mhc_mult]
        norm_weight: optional RMSNorm weight of shape [mhc_mult * hidden_size]
        norm_eps: epsilon for RMSNorm
        mhc_mult: number of hyper-connection heads (currently only 4 is guaranteed to work)
        pre_eps: epsilon for pre-layer mix sigmoid
        n_splits: number of splits for the split-K GEMM; the platform default applies

    Returns:
        layer_input: tensor of shape [..., hidden_size], input to lm_head
    """
    n_splits = default_n_splits() if n_splits is None else n_splits
    mhc_mult3 = mhc_mult * (2 + mhc_mult)

    if fn.shape[0] < mhc_mult3:
        fn = F.pad(fn, (0, 0, 0, mhc_mult3 - fn.shape[0]))

    mixes = mhc_pre_norm_fn(residual, fn, norm_weight, norm_eps, n_splits=n_splits)

    mixes = mixes[..., :mhc_mult].contiguous()

    if scale.numel() == 1:
        scale = scale.reshape(1)
    mix = mhc_head_compute_mix(mixes, scale, base, pre_eps)

    return mhc_pre_apply_mix(residual, mix.unsqueeze(-1))


def mhc_pre_big_fuse(
    residual: torch.Tensor,
    fn: torch.Tensor,
    mhc_scale: torch.Tensor,
    mhc_base: torch.Tensor,
    rms_eps: float,
    mhc_pre_eps: float,
    mhc_sinkhorn_eps: float,
    mhc_post_mult_value: float,
    sinkhorn_repeat: int,
    n_splits: int | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Fused inference path: prenorm GEMM + reduce/rmsnorm + split mixes + sinkhorn + apply mix.

    Only used with gradients disabled: the GEMM is called directly and the epilogue stays
    in the tile-kernels operator. The FP32 fn matrix already includes the norm weight.
    """
    n_splits = default_n_splits() if n_splits is None else n_splits
    assert residual.dtype == torch.bfloat16
    assert fn.dtype == torch.float32
    assert mhc_scale.dtype == torch.float32
    assert mhc_base.dtype == torch.float32

    mhc_mult = residual.shape[-2]
    hidden_size = residual.shape[-1]
    mhc_mult2 = mhc_mult * mhc_mult
    mhc_mult3 = mhc_mult * 2 + mhc_mult2

    mhc_hidden_size = mhc_mult * hidden_size
    assert fn.shape[0] == mhc_mult3
    assert fn.shape[1] == mhc_hidden_size
    assert mhc_scale.shape == (3,)
    assert mhc_base.shape == (mhc_mult3,)

    outer_shape = residual.shape[:-2]

    residual_flat = residual.view(-1, mhc_mult, hidden_size)
    num_tokens = residual_flat.shape[0]

    post_mix = torch.empty(num_tokens, mhc_mult, dtype=torch.float32, device=residual.device)
    comb_mix = torch.empty(num_tokens, mhc_mult2, dtype=torch.float32, device=residual.device)
    layer_input = torch.empty(num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device)

    gemm_out_mul, gemm_out_sqrsum = mhc_prenorm_gemm(residual_flat, fn, n_splits)

    actual_n_splits = gemm_out_mul.shape[0]
    gemm_out_mul = gemm_out_mul.view(actual_n_splits, num_tokens, mhc_mult3)
    gemm_out_sqrsum = gemm_out_sqrsum.view(actual_n_splits, num_tokens)

    mhc_pre_big_fuse_fwd(
        gemm_out_mul,
        gemm_out_sqrsum,
        mhc_scale,
        mhc_base,
        residual_flat,
        post_mix,
        comb_mix,
        layer_input,
        rms_eps,
        mhc_pre_eps,
        mhc_sinkhorn_eps,
        mhc_post_mult_value,
        sinkhorn_repeat,
        n_splits=n_splits,
    )

    post_mix = post_mix.view(*outer_shape, mhc_mult, 1)
    comb_mix = comb_mix.view(*outer_shape, mhc_mult, mhc_mult)
    layer_input = layer_input.view(*outer_shape, hidden_size)

    return post_mix, comb_mix, layer_input