"""SSL-Mamba network (inference-only port).

Source: github.com/hoanmyTran/spoofing_ssl_mamba @ 9afc33de7f74a9b9b3423a9028c10e1d18c1c6b8 (MIT),
``model.py`` ``SSLClassifier``. Checkpoint: HF hoanmyTran/ssl_spoofing_mamba @
4e72dfcef4f03460e08f68a4a5fbcf3705a3414e ``interspeech2025.pt`` (Apache-2.0; flat state_dict),
converted once to ``model.safetensors`` with the training-only ``loss_fn.weight`` dropped.

Edits vs the source (all inference-only; module names unchanged so keys load strictly):
  1. pytorch_lightning / torchmetrics scaffolding, loss and training/eval loops removed.
  2. XLS-R 300M built from the bundled config (``xlsr_300m_config.json`` = the
     facebook/wav2vec2-xls-r-300m config) instead of ``from_pretrained``: no Hub weight
     download; the checkpoint overwrites every weight (the front-end is fully fine-tuned).
  3. ``mamba_ssm.Mamba`` -> ``_mamba_v1.Mamba`` (pure PyTorch, same parameter names, same
     math; chunked closed-form selective scan, see that file). mamba_ssm / causal_conv1d are
     not installed and Triton is unreliable on sm_120.
  4. ``multiconv_cgmlp`` espnet imports -> ``_espnet_shim`` (LayerNorm eps 1e-12, activations);
     ``pooling.py`` (MultiHeadAttentionPooling) vendored from the authors' sibling repo
     (dissimilarity_deepfake_detection, as in benchmarks/SSL-Feature-Gating/), not in this repo.
  5. ``nn.SiLU(inplace=True)`` / ``nn.SELU(inplace=True)`` -> non-inplace (export-friendly;
     identical values).
Forward is otherwise verbatim: 25 hidden states -> Linear(1024,128) -> SiLU -> sum over layers
-> MultiConv cgMLP (7/15/23/31) -> 12 x [x = x + LayerNorm(Mamba(x))] (unidirectional)
-> MHAP (4 heads, mean+std) -> Linear 256->128 -> SELU -> Linear 128->2.
"""
from __future__ import annotations

import os

import numpy as np
import torch
import torch.nn as nn
from transformers import Wav2Vec2Config, Wav2Vec2Model

from _mamba_v1 import Mamba
from multiconv_cgmlp import MultiConvolutionalGatingMLP
from pooling import MultiHeadAttentionPooling

_DIR = os.path.dirname(os.path.abspath(__file__))
_CFG = os.path.join(_DIR, "xlsr_300m_config.json")
WINDOW = 64600  # source config dataset.max_len (4.0375 s @ 16 kHz)


def pad_repeat(x: np.ndarray, max_len: int = WINDOW) -> np.ndarray:
    """Source dataset.pad: first max_len samples; shorter clips tile-repeated."""
    x = np.asarray(x, dtype=np.float32).reshape(-1)
    n = x.shape[0]
    if n >= max_len:
        return x[:max_len]
    return np.tile(x, max_len // n + 1)[:max_len].astype(np.float32)


class SSLMambaNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.feature_extractor = Wav2Vec2Model(Wav2Vec2Config.from_json_file(_CFG))
        self.cgmlp = MultiConvolutionalGatingMLP(
            size=128, linear_units=1024, arch_type="concat_fusion", kernel_sizes="7,15,23,31",
            merge_conv_kernel=31, use_non_linear=True, dropout_rate=0.1,
            use_linear_after_conv=True, activation="silu", gate_activation="silu")
        self.feature_projection = nn.Linear(1024, 128)
        self.blocks = nn.ModuleList(
            nn.Sequential(Mamba(d_model=128, d_state=16, d_conv=4, expand=2), nn.LayerNorm(128))
            for _ in range(12))
        self.last_pooling = MultiHeadAttentionPooling(128)
        self.silu = nn.SiLU()
        self.classifier = nn.Sequential(nn.Linear(256, 128), nn.SELU(), nn.Linear(128, 2))

    def set_ref_scan(self, on: bool) -> None:
        for m in self.modules():
            if isinstance(m, Mamba):
                m.use_ref_scan = on

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.feature_extractor(x, output_hidden_states=True)
        hs = torch.stack(x.hidden_states, dim=1)          # (B, 25, T, 1024)
        hs = self.silu(self.feature_projection(hs))
        x = torch.sum(hs, 1)                               # (B, T, 128)
        x = self.cgmlp(x)
        for b in self.blocks:
            x = x + b(x)
        x = self.last_pooling(x.permute(0, 2, 1))          # (B, 256, 1)
        return self.classifier(x.squeeze(-1))              # (B, 2); class 1 = bona fide


def load_state(path: str) -> dict:
    from safetensors.torch import load_file
    return load_file(path)