"""Automatic selection and inference wrapper for the bundled Mamba-3 MIMO kernels.

Includes Mamba-3 forward integration adapted by Continuous Intelligence:
https://github.com/state-spaces/mamba/tree/a14b1dff0454a3bc27d9eb31355dc01e4b2490ec

Copyright (c) 2026, Dao AI Lab, Goombalab.
Upstream portions are licensed under Apache-2.0; see NOTICE.
"""

from __future__ import annotations

import warnings

import torch
import torch.nn.functional as F
from einops import rearrange
from torch import Tensor, nn

from ._mamba import Mamba3


def configure_backend(model: nn.Module) -> None:
    """Enable fused inference only for supported hardware and cell geometry."""
    model.backend = "torch"
    device = next(model.parameters()).device
    if device.type != "cuda" or torch.version.hip is not None:
        return
    cells = [cell for cell in model.modules() if isinstance(cell, Mamba3)]
    if torch.cuda.get_device_capability(device)[0] < 8 or not cells:
        return
    # Float32 tiles need more shared memory than bf16. Use smaller chunks on
    # Ampere and consumer GPUs.
    chunk_budget = (
        64
        if torch.cuda.get_device_properties(device).shared_memory_per_block_optin >= 196608
        else 32
    )
    if any(
        not cell.is_mimo
        or cell.mimo_rank not in (1, 2, 4, 8)
        or cell.d_state not in (32, 64, 128)
        or cell.headdim not in (32, 64, 128)
        or max(8, chunk_budget // cell.mimo_rank) * cell.mimo_rank > chunk_budget
        for cell in cells
    ):
        return
    try:
        from .._vendor.mamba3.mamba3_mimo_fwd import mamba_mimo_forward  # noqa: F401
    except (ImportError, OSError) as error:
        warnings.warn(
            f"Fused Mamba-3 is unavailable; using portable PyTorch on CUDA: {error}",
            RuntimeWarning,
            stacklevel=2,
        )
        return
    for cell in cells:
        cell.chunk_size = max(8, chunk_budget // cell.mimo_rank)
        cell.fused = True
    model.backend = "mamba3"


def mamba3_forward(cell: Mamba3, u: Tensor) -> Tensor:
    """Run the forward path, including rank-scaled chunk padding."""
    from .._vendor.mamba3.angle_dt import angle_dt_fwd
    from .._vendor.mamba3.layernorm_gated import _layer_norm_fwd
    from .._vendor.mamba3.mamba3_mimo_fwd import mamba_mimo_forward
    from .._vendor.mamba3.mamba3_mimo_utils import compute_dacs_segsum_triton

    seq = u.shape[1]
    chunk = cell.chunk_size
    u = F.pad(u, (0, 0, 0, (-seq) % chunk))
    bc = cell.d_state * cell.mimo_rank
    z, x, B, C, dd_dt, dd_A, trap, angles = torch.split(
        cell.in_proj(u),
        [
            cell.d_inner,
            cell.d_inner,
            bc,
            bc,
            cell.nheads,
            cell.nheads,
            cell.nheads,
            cell.num_rope_angles,
        ],
        dim=-1,
    )
    x = rearrange(x, "b l (h p) -> b l h p", p=cell.headdim)
    z = rearrange(z, "b l (h p) -> b l h p", p=cell.headdim)
    B = rearrange(B, "b l (r g n) -> b l r g n", r=cell.mimo_rank, g=1)
    C = rearrange(C, "b l (r g n) -> b l r g n", r=cell.mimo_rank, g=1)
    # Keep the kernel's reduction and rounding order, including 1 / sqrt.
    B = _layer_norm_fwd(
        B.reshape(-1, cell.d_state),
        cell.B_norm.weight,
        None,
        cell.B_norm.eps,
        is_rms_norm=True,
    )[0].view_as(B)
    C = _layer_norm_fwd(
        C.reshape(-1, cell.d_state),
        cell.C_norm.weight,
        None,
        cell.C_norm.eps,
        is_rms_norm=True,
    )[0].view_as(C)
    A = (-F.softplus(dd_A.float())).clamp(max=-1e-4)
    dt = F.softplus(dd_dt + cell.dt_bias)
    adt = rearrange(A * dt, "b l h -> b h l").contiguous()
    dt = rearrange(dt, "b l h -> b h l").contiguous()
    angles = angles[:, :, None].expand(-1, -1, cell.nheads, -1).float().contiguous()
    with torch.cuda.device(u.device):
        cumulative, _ = angle_dt_fwd(angles, dt, chunk_size=chunk, return_output_state=True)
        da_cs, da_cs_rev, segsum = compute_dacs_segsum_triton(adt, chunk)
        y, _, _ = mamba_mimo_forward(
            C.contiguous(),
            B.contiguous(),
            x.contiguous(),
            cell.C_bias,
            cell.B_bias,
            cell.mimo_x,
            cell.mimo_o,
            z.contiguous(),
            cell.D,
            cell.mimo_z,
            cumulative,
            da_cs,
            da_cs_rev,
            dt,
            rearrange(trap, "b l h -> b h l").contiguous(),
            segsum,
            chunk_size=chunk,
            rotary_dim_divisor=4,
            dtype=x.dtype,
        )
    return cell.out_proj(rearrange(y, "b l h p -> b l (h p)").to(u.dtype))[:, :seq]