"""Triton-3.2-compatible Mamba3 SISO wrapper.

The upstream Mamba3 SISO Triton kernel in the release graft uses
``tl.make_tensor_descriptor``. HF A10 Jobs require Triton 3.2.0 for working
driver discovery with torch 2.6, and Triton 3.2.0 does not expose that API.

This module keeps the public ``mamba3_siso_combined`` entrypoint intact and
uses regular Torch ops for the SISO path. It is intentionally conservative:
no silent CPU fallback, no fake outputs, and no dependency on newer Triton APIs.
"""

from __future__ import annotations

from typing import Optional, Tuple

import torch
from torch import Tensor
import torch.nn.functional as F


def _expand_qk_heads(x: Tensor, nheads: int) -> Tensor:
    """Expand grouped Q/K heads to value-head count."""
    nheads_qk = x.shape[2]
    if nheads_qk == nheads:
        return x
    if nheads % nheads_qk != 0:
        raise ValueError(f"nheads={nheads} must be divisible by nheads_qk={nheads_qk}")
    repeat = nheads // nheads_qk
    return x.repeat_interleave(repeat, dim=2)


def _rotate_pairwise(x: Tensor, angles: Tensor) -> Tensor:
    """Apply adjacent-pair RoPE to the first 2 * angles.shape[-1] channels."""
    rotary_pairs = angles.shape[-1]
    rotary_dim = rotary_pairs * 2
    if rotary_dim == 0:
        return x
    if x.shape[-1] < rotary_dim:
        raise ValueError(f"rotary_dim={rotary_dim} exceeds q/k dim={x.shape[-1]}")

    x_rot = x.clone()
    x0 = x[..., :rotary_dim:2]
    x1 = x[..., 1:rotary_dim:2]
    cos = torch.cos(angles.to(torch.float32)).to(dtype=x.dtype)
    sin = torch.sin(angles.to(torch.float32)).to(dtype=x.dtype)
    x_rot[..., :rotary_dim:2] = x0 * cos - x1 * sin
    x_rot[..., 1:rotary_dim:2] = x0 * sin + x1 * cos
    return x_rot


def _torch_siso_scan(
    Q: Tensor,
    K: Tensor,
    V: Tensor,
    ADT: Tensor,
    DT: Tensor,
    Trap: Tensor,
    Q_bias: Tensor,
    K_bias: Tensor,
    Angles: Tensor,
    D: Optional[Tensor],
    Z: Optional[Tensor],
    Input_States: Optional[Tuple[Tensor, Tensor, Tensor]],
    return_final_states: bool,
) -> Tensor | Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
    if not (Q.is_cuda and K.is_cuda and V.is_cuda):
        raise RuntimeError("Mamba3 Torch SISO fallback requires CUDA tensors")

    batch, seqlen, _, qk_dim = Q.shape
    _, _, nheads, headdim_v = V.shape

    q = _expand_qk_heads(Q, nheads)
    k = _expand_qk_heads(K, nheads)

    q = q + Q_bias.to(dtype=q.dtype, device=q.device).view(1, 1, nheads, qk_dim)
    k = k + K_bias.to(dtype=k.dtype, device=k.device).view(1, 1, nheads, qk_dim)

    # Upstream angle_dt computes cumulative angle rates weighted by dt.
    dt_seq = DT.transpose(1, 2).to(torch.float32)  # (B, L, H)
    angle_cumsum = torch.cumsum(Angles.to(torch.float32) * dt_seq.unsqueeze(-1), dim=1)
    q = _rotate_pairwise(q, angle_cumsum)
    k_unscaled = _rotate_pairwise(k, angle_cumsum)

    trap = torch.sigmoid(Trap.to(torch.float32))
    dt = DT.to(torch.float32)
    shifted_dt = torch.cat([dt[:, :, 1:], torch.zeros_like(dt[:, :, :1])], dim=2)
    shifted_trap = torch.cat([trap[:, :, 1:], torch.zeros_like(trap[:, :, :1])], dim=2)
    gamma = dt * trap
    scale = shifted_dt * (1.0 - shifted_trap) + gamma
    k_scaled = k_unscaled * scale.transpose(1, 2).unsqueeze(-1).to(k_unscaled.dtype)

    if Input_States is not None:
        state, k_state, v_state = Input_States
        state_acc = state.to(torch.float32).clone()
        first_dt = dt[:, :, 0].unsqueeze(-1).unsqueeze(-1)
        first_trap = trap[:, :, 0].unsqueeze(-1).unsqueeze(-1)
        state_acc = state_acc + v_state.to(torch.float32).unsqueeze(-1) * k_state.to(torch.float32).unsqueeze(-2) * first_dt * (1.0 - first_trap)
    else:
        state_acc = torch.zeros(batch, nheads, headdim_v, qk_dim, device=V.device, dtype=torch.float32)

    d = torch.zeros(nheads, device=V.device, dtype=torch.float32) if D is None else D.to(torch.float32)
    outputs: list[Tensor] = []

    for t in range(seqlen):
        q_t = q[:, t].to(torch.float32)  # (B, H, N)
        k_t = k_scaled[:, t].to(torch.float32)
        k_raw_t = k_unscaled[:, t].to(torch.float32)
        v_t = V[:, t].to(torch.float32)  # (B, H, P)

        decayed = state_acc * torch.exp(ADT[:, :, t].to(torch.float32)).unsqueeze(-1).unsqueeze(-1)
        state_out = torch.einsum("bhn,bhpn->bhp", q_t, decayed)
        qk_dot = (q_t * k_raw_t).sum(dim=-1) * gamma[:, :, t]
        y_t = state_out + (d.view(1, nheads) + qk_dot).unsqueeze(-1) * v_t

        if Z is not None:
            y_t = y_t * F.silu(Z[:, t].to(torch.float32))
        outputs.append(y_t.to(V.dtype))

        state_acc = decayed + torch.einsum("bhp,bhn->bhpn", v_t, k_t)

    out = torch.stack(outputs, dim=1)
    if not return_final_states:
        return out

    final_angle = angle_cumsum[:, -1]
    final_k = k_unscaled[:, -1]
    final_v = V[:, -1]
    return out, final_angle, state_acc, final_k, final_v


def mamba3_siso_combined(
    Q: Tensor,
    K: Tensor,
    V: Tensor,
    ADT: Tensor,
    DT: Tensor,
    Trap: Tensor,
    Q_bias: Tensor,
    K_bias: Tensor,
    Angles: Tensor,
    D: Optional[Tensor] = None,
    Z: Optional[Tensor] = None,
    Input_States: Optional[Tuple[Tensor, Tensor, Tensor, Tensor]] = None,
    chunk_size: int = 64,
    return_final_states: bool = False,
    cu_seqlens: Optional[Tensor] = None,
) -> Tensor | Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
    del chunk_size
    if cu_seqlens is not None:
        raise NotImplementedError("Torch Mamba3 SISO fallback does not support varlen cu_seqlens")

    angle_state, ssm_state, k_state, v_state = (
        Input_States if Input_States is not None else (None, None, None, None)
    )
    if any(x is not None for x in (angle_state, ssm_state, k_state, v_state)):
        if not all(x is not None for x in (angle_state, ssm_state, k_state, v_state)):
            raise ValueError("Mamba3 input states must be all present or all absent")
        torch_states = (ssm_state, k_state, v_state)
    else:
        torch_states = None

    return _torch_siso_scan(
        Q=Q,
        K=K,
        V=V,
        ADT=ADT,
        DT=DT,
        Trap=Trap,
        Q_bias=Q_bias,
        K_bias=K_bias,
        Angles=Angles,
        D=D,
        Z=Z,
        Input_States=torch_states,
        return_final_states=return_final_states,
    )