"""XLSR-MamBo network (MamBo-3: Hydra-N3, 5 repeats), ported from
https://github.com/saki-ciallo/XLSR-MamBo (MIT) @ f429bb1:

- model_build.F4_XLSR300M_General            -> XLSRMamBo
- modules/step1_feature_ptm.WLSR300M_Extractor + RMS_InputProj
- modules/step2_model_blocks.{Hydra_Block, MLP_block, MHA_Block, MLP_SwiGLU,
  Mix_n_Hydra_1MLP_1MHA_1MLP}
- modules/step3_pooling_outputs.GatedAttentionPool
- modules/hydra (Hydra mixer, vendored in ./_hydra)
- train_config: task xlsr_Hydra = Mamba2Config(d_model=128, d_state=64, d_conv=7,
  expand=2, mha_num_heads=4, mha_mlp_dim=0, mha_d_conv=4, residual_in_fp32=True,
  fused_add_norm=True, use_mem_eff_path=True); block_layers=5, n_mamba=3.

Parameter/module names are kept identical to the source so the released
state_dict loads strictly.

Two numerically-equivalent backends:
  * ``pure_torch=False`` (default): the source path -- Triton kernels from the
    vendored mamba_ssm 2.2.6.post3 (fused add+RMSNorm, gated RMSNorm, Mamba2 SSD
    chunk scan via Hydra's mem-efficient path).
  * ``pure_torch=True``: the same math in plain PyTorch (explicit quadratic SSD
    form, plain RMSNorm). Exportable to ONNX/TensorRT and used as a fallback on
    GPUs where the Triton kernels misbehave. Parity is checked in the tests.

Deviation from the source (deliberate): torchaudio's bundle wrapper applies
``layer_norm(waveforms, waveforms.shape)`` -- jointly over the whole (B, T)
batch, so the source's scores depend on batch composition. We normalise each
utterance on its own (== the source at batch size 1), making scores
batch-independent and reproducible.
"""
from __future__ import annotations

import os
import sys
from dataclasses import dataclass
from typing import Optional, Tuple

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

_HERE = os.path.dirname(os.path.abspath(__file__))
if _HERE not in sys.path:
    sys.path.insert(0, _HERE)  # vendored mamba_ssm + _hydra resolve from here

from mamba_ssm.modules.mha import MHA  # noqa: E402
from mamba_ssm.ops.triton.layer_norm import RMSNorm, rms_norm_fn  # noqa: E402
from _hydra import Hydra  # noqa: E402

WINDOW = 66800  # source default --fixed-length-samples (4.175 s @ 16 kHz)


@dataclass
class MamBoConfig:
    d_model: int = 128
    d_state: int = 64
    d_conv: int = 7          # xlsr_Hydra task: MODEL_CONFIG.replace(d_conv=7)
    expand: int = 2
    mha_num_heads: int = 4
    mha_mlp_dim: int = 0
    mha_d_conv: int = 4
    residual_in_fp32: bool = True
    use_mem_eff_path: bool = True
    block_layers: int = 5
    n_mamba: int = 3


# ----------------------------------------------------------------------------
# pure-torch equivalents of the Triton ops
# ----------------------------------------------------------------------------
def _rms(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
    xf = x.float()
    rstd = torch.rsqrt(xf.square().mean(dim=-1, keepdim=True) + eps)
    return (xf * rstd * weight.float()).to(x.dtype)


def add_rmsnorm(norm: RMSNorm, hidden, residual, prenorm: bool, residual_in_fp32: bool, pure: bool):
    """mamba_ssm fused add + RMSNorm (rms_norm_fn) semantics."""
    if not pure:
        return rms_norm_fn(hidden, norm.weight, norm.bias, residual=residual, prenorm=prenorm,
                           residual_in_fp32=residual_in_fp32, eps=norm.eps)
    res = hidden if residual is None else hidden + residual
    if residual_in_fp32:
        res = res.float()
    out = _rms(res, norm.weight, norm.eps).to(hidden.dtype)
    return (out, res) if prenorm else out


def ssd_scan_torch(x, dt, A, B, C):
    """Mamba2 SSD scan, exact quadratic form (no chunking; L is ~208 here).

    x: (b, l, h, p)  dt: (b, l, h) (already softplus'ed)  A: (h,)
    B, C: (b, l, g, n) with g == 1.  Returns y: (b, l, h, p), fp32 math.
    y_t = sum_{s<=t} (C_t . B_s) * exp(sum_{r=s+1..t} A dt_r) * dt_s * x_s
    """
    dtype = x.dtype
    x, dt, B, C = x.float(), dt.float(), B.float(), C.float()
    dA = dt * A.float()[None, None, :]                       # (b, l, h)
    cs = torch.cumsum(dA, dim=1)                              # (b, l, h)
    seg = cs[:, :, None, :] - cs[:, None, :, :]               # (b, t, s, h)
    L = x.shape[1]
    causal = torch.ones(L, L, dtype=torch.bool, device=x.device).tril()
    decay = torch.exp(seg.masked_fill(~causal[None, :, :, None], float("-inf")))
    CB = torch.einsum("btgn,bsgn->bts", C, B)                 # g == 1
    w = CB[..., None] * decay * dt[:, None, :, :]             # (b, t, s, h)
    y = torch.einsum("btsh,bshp->bthp", w, x)
    return y.to(dtype)


def _chunk_flip_join(t, dim, op):
    h = t.shape[dim] // 2  # static under tracing; slicing (not chunk) -> no ONNX Sequence ops
    fwd, bwd = t.narrow(dim, 0, h), t.narrow(dim, h, h)
    bwd = torch.flip(bwd, dims=(1,))
    if op == "sum":
        return fwd + bwd
    if op == "vstack":
        return torch.cat([fwd, bwd], dim=0)
    raise ValueError(op)


def hydra_forward_torch(m: Hydra, u: torch.Tensor) -> torch.Tensor:
    """Pure-torch mirror of _hydra/ops.py HydraSplitConv1dScanCombinedFn.forward
    (the use_mem_eff_path=True path the checkpoint was trained/evaluated with)."""
    batch, seqlen, _ = u.shape
    d_inner, d_state, headdim, ngroups = m.d_inner, m.d_state, m.headdim, m.ngroups
    nheads = d_inner // headdim
    d_xBC = d_inner + 2 * (2 * ngroups * d_state)
    zxbcdt = m.in_proj(u)
    A = -torch.exp(m.A_log.float())
    z = zxbcdt[..., :d_inner]
    xBC_og = zxbcdt[..., d_inner:d_inner + d_xBC]
    dt = zxbcdt[..., d_inner + d_xBC:]
    k = int(m.d_conv)  # odd (7): "same" padding == k // 2 each side (python int: traceable)
    xBC = F.silu(F.conv1d(xBC_og.transpose(1, 2), m.conv1d.weight, m.conv1d.bias, padding=k // 2,
                          groups=int(m.conv1d.groups)).transpose(1, 2))
    x, BC = xBC[..., :d_inner], xBC[..., d_inner:]
    x_og = x
    x = torch.cat([x, torch.flip(x, dims=(1,))], dim=0)
    BC = _chunk_flip_join(BC, dim=-1, op="vstack")
    nBC = ngroups * d_state
    Bm, Cm = BC[..., :nBC], BC[..., nBC:]
    x = rearrange(x, "b l (h p) -> b l h p", p=headdim)
    Bm = rearrange(Bm, "b l (g n) -> b l g n", g=ngroups)
    Cm = rearrange(Cm, "b l (g n) -> b l g n", g=ngroups)
    D = repeat(F.linear(x_og, m.fc_D.weight, m.D), "b l h -> b l (h p)", p=headdim)
    dt = _chunk_flip_join(dt, dim=-1, op="vstack")
    dt = F.softplus(dt.float() + m.dt_bias.float())
    scan = ssd_scan_torch(x, dt, A, Bm, Cm)
    scan = rearrange(scan, "b l h p -> b l (h p)")
    y = torch.roll(scan, shifts=1, dims=1)
    y = torch.cat([torch.zeros_like(y[:, :1]), y[:, 1:]], dim=1)  # y[:, 0] = 0 (export-friendly)
    y = _chunk_flip_join(y, dim=0, op="sum") + D * x_og
    # gated RMSNorm, norm_before_gate=True
    u_ = _rms(y, m.norm.weight, m.norm.eps).float() * F.silu(z.float())
    return F.linear(u_.to(y.dtype), m.out_proj.weight, m.out_proj.bias)


# ----------------------------------------------------------------------------
# blocks (Add -> RMSNorm -> mixer), names as in step2_model_blocks.py
# ----------------------------------------------------------------------------
class MLP_SwiGLU(nn.Module):
    def __init__(self, d_model: int):
        super().__init__()
        hidden_dim = int(8 * d_model / 3)
        self.w1_w3_combined = nn.Linear(d_model, 2 * hidden_dim, bias=False)
        self.w2_out = nn.Linear(hidden_dim, d_model, bias=False)
        self.activation = nn.SiLU()

    def forward(self, x):
        gv = self.w1_w3_combined(x)
        h = self.w2_out.in_features
        gate, value = gv[..., :h], gv[..., h:]
        return self.w2_out(self.activation(gate) * value)  # dropout(0.1) is a no-op at eval


class _PreNormBlock(nn.Module):
    pure = False

    def __init__(self, cfg: MamBoConfig):
        super().__init__()
        self.norm = RMSNorm(cfg.d_model)
        self.residual_in_fp32 = cfg.residual_in_fp32

    def mix(self, h):
        raise NotImplementedError

    def forward(self, hidden, residual=None):
        hidden, residual = add_rmsnorm(self.norm, hidden, residual, True, self.residual_in_fp32, self.pure)
        return self.mix(hidden), residual


class Hydra_Block(_PreNormBlock):
    def __init__(self, cfg: MamBoConfig):
        super().__init__(cfg)
        # NB: source passes no headdim -> Hydra default 64 (d_inner 256, 4 heads)
        self.mamba = Hydra(d_model=cfg.d_model, d_state=cfg.d_state, d_conv=cfg.d_conv,
                           expand=cfg.expand, use_mem_eff_path=cfg.use_mem_eff_path)

    def mix(self, h):
        return hydra_forward_torch(self.mamba, h) if self.pure else self.mamba(h)


class MLP_block(_PreNormBlock):
    def __init__(self, cfg: MamBoConfig):
        super().__init__(cfg)
        self.mlp = MLP_SwiGLU(cfg.d_model)

    def mix(self, h):
        return self.mlp(h)


class MHA_Block(_PreNormBlock):
    def __init__(self, cfg: MamBoConfig):
        super().__init__(cfg)
        self.mha = MHA(embed_dim=cfg.d_model, num_heads=cfg.mha_num_heads,
                       mlp_dim=cfg.mha_mlp_dim, d_conv=cfg.mha_d_conv)

    def mix(self, h):
        # MHA without flash_attn/causal_conv1d: nn.Conv1d (causal trim) + SDPA (non-causal)
        return self.mha(h)


class Mix_n_Hydra_1MLP_1MHA_1MLP(nn.Module):
    def __init__(self, cfg: MamBoConfig):
        super().__init__()
        self.layers = nn.ModuleList()
        for _ in range(cfg.block_layers):
            self.layers.extend([Hydra_Block(cfg) for _ in range(cfg.n_mamba)])
            self.layers.append(MLP_block(cfg))
            self.layers.append(MHA_Block(cfg))
            self.layers.append(MLP_block(cfg))
        self.norm = RMSNorm(cfg.d_model)
        self.residual_in_fp32 = cfg.residual_in_fp32
        self.pure = False

    def forward(self, hidden):
        residual = None
        for layer in self.layers:
            hidden, residual = layer(hidden, residual)
        return add_rmsnorm(self.norm, hidden, residual, False, self.residual_in_fp32, self.pure)


class StackedMamba2(nn.Module):
    def __init__(self, cfg: MamBoConfig):
        super().__init__()
        self.layers = Mix_n_Hydra_1MLP_1MHA_1MLP(cfg)

    def forward(self, x):
        return self.layers(x)


class RMS_InputProj(nn.Module):
    def __init__(self, d_model: int):
        super().__init__()
        self.input_proj = nn.Linear(1024, d_model, bias=False)
        self.norm = RMSNorm(d_model)
        self.pure = False

    def forward(self, x):
        x = self.input_proj(x)
        x = _rms(x, self.norm.weight, self.norm.eps) if self.pure else self.norm(x)
        return F.silu(x)


class GatedAttentionPool(nn.Module):
    def __init__(self, d_model: int):
        super().__init__()
        hidden_dim = int(8 * d_model // 3)
        self.gate_value_proj = nn.Linear(d_model, 2 * hidden_dim, bias=False)
        self.to_score = nn.Linear(hidden_dim, 1, bias=False)
        self.activation = nn.SiLU()

    def forward(self, x):
        gv = self.gate_value_proj(x)
        h = self.to_score.in_features
        gate, value = gv[..., :h], gv[..., h:]
        scores = self.to_score(self.activation(gate) * value).squeeze(-1)  # (B, L)
        weights = F.softmax(scores, dim=1).unsqueeze(1)                     # (B, 1, L)
        return torch.bmm(weights, x).squeeze(1)                            # (B, D)


class _BundleWrapper(nn.Module):
    """Key-compatible stand-in for torchaudio's pipelines `_Wav2Vec2Model` wrapper
    (keys `model.*`), but with PER-UTTERANCE waveform normalisation."""

    def __init__(self):
        super().__init__()
        import torchaudio
        # architecture only; the fine-tuned checkpoint overwrites every weight
        self.model = torchaudio.models.wav2vec2_xlsr_300m()

    def forward(self, wav):
        # per-utterance layer_norm (see module docstring), written out explicitly
        # (== F.layer_norm(wav, wav.shape[-1:]), eps 1e-5) so it exports to ONNX.
        mean = wav.mean(dim=-1, keepdim=True)
        var = (wav - mean).square().mean(dim=-1, keepdim=True)
        wav = (wav - mean) * torch.rsqrt(var + 1e-5)
        out, _ = self.model(wav)
        return out


class WLSR300M_Extractor(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = _BundleWrapper()

    def forward(self, wav):
        return self.model(wav)  # (B, L, 1024), last transformer layer


class XLSRMamBo(nn.Module):
    def __init__(self, cfg: MamBoConfig | None = None, pure_torch: bool = False):
        super().__init__()
        cfg = cfg or MamBoConfig()
        self.feature_extractor = WLSR300M_Extractor()
        self.input_proj = RMS_InputProj(cfg.d_model)
        self.stacked_mamba2 = StackedMamba2(cfg)
        self.attention_pool = GatedAttentionPool(cfg.d_model)
        self.classifier = nn.Linear(cfg.d_model, 2)
        self.set_pure_torch(pure_torch)

    def set_pure_torch(self, flag: bool) -> None:
        for mod in self.modules():
            if hasattr(mod, "pure"):
                mod.pure = flag

    def forward(self, wav: torch.Tensor) -> torch.Tensor:
        x = self.feature_extractor(wav)
        x = self.input_proj(x)
        x = self.stacked_mamba2(x)
        x = self.attention_pool(x)
        return self.classifier(x)  # (B, 2) logits; class 1 = bona fide


def load_state(path: str) -> dict:
    sd = torch.load(path, map_location="cpu", weights_only=True)
    return {(k[len("_orig_mod."):] if k.startswith("_orig_mod.") else k): v for k, v in sd.items()}


def pad_repeat(wav, max_len: int = WINDOW):
    """Source pad_wave_repeat: first max_len samples; shorter clips are tiled."""
    import numpy as np
    n = len(wav)
    if n == 0:
        raise ValueError("empty waveform")
    if n >= max_len:
        return wav[:max_len]
    return np.tile(wav, -(-max_len // n))[:max_len]