"""Fake-Mamba network (XLS-R 300M + PN-BiMamba encoder), ported from
https://github.com/xuanxixi/Fake-Mamba (no LICENSE upstream) @ e966d2f:

- model.py:SSLModel                    -> SSLModel
- model.py:Fake_Mamba                  -> FakeMambaNet
- conformer00.py:PN_BiMambas_Encoder   -> PN_BiMambas_Encoder
- conformer00.py:PN_BiMambas           -> PN_BiMambas
- mamba_ssm 1.x modules/mamba_simple.Mamba + ops/selective_scan_interface
  .selective_scan_ref (vendored subset, as in benchmarks/RawBMamba)  -> Mamba, selective_scan_ref

Parameter / module names are identical to the source so the released Exp65
``best.pth`` loads with strict=True.

Edits vs. the source (none changes the math):
  1. model.py imports two modules that are not in the repo
     (``conformer.ConformerBlock, BiMambaEncoder``; conformer00 imports
     ``ConExBiMamba.ExBiMamba``). They are used only by other back-ends, so only
     Fake_Mamba + PN_BiMambas are ported; nothing is stubbed.
  2. ``PN_BiMambas(dim=144, depth=12)``: the released checkpoint Exp65 has 12
     PN-BiMamba blocks (encoder.layers.0-11); the source/paper default is 7.
  3. SSLModel builds the wav2vec2 XLS-R 300M architecture from its stored model
     config (``xlsr2_300m_cfg.json``, extracted from fairseq's xlsr2_300m.pt and
     merged onto the installed Wav2Vec2Config -- the cross-version fix used in
     W2V2-AASIST) instead of ``load_model_ensemble_and_task`` on a hard-coded
     path; the fine-tuned checkpoint overwrites every weight. ``extract_feat``
     no longer calls ``self.model.train()`` on a device/dtype mismatch (it would
     re-enable dropout at inference); the wrapper moves the net to the device
     before ``.eval()`` anyway.
  4. Mamba: mamba_ssm's CUDA fast path (``mamba_inner_fn`` / causal_conv1d /
     selective_scan_cuda) is replaced by the library's own pure-PyTorch
     reference (``use_fast_path=False`` branch: nn.Conv1d + ``selective_scan_ref``),
     identical math; runs on every GPU (incl. sm_120) and exports to ONNX/TensorRT.
     The inference-cache (step) path is dropped (unused). The time recurrence
     is evaluated with a Hillis-Steele parallel prefix scan
     (``selective_scan_parallel``; same recurrence, float-rounding-level
     differences, ~log2(L) tensor ops instead of an L-step loop);
     ``FAKE_MAMBA_SCAN=ref`` selects the library's sequential loop (parity is
     checked in the tests).
  5. PN_BiMambas_Encoder runs the forward and the flipped (backward) sequence
     through the shared Mamba in ONE call (concatenated along the batch axis)
     instead of two calls -- same weights, same per-sample math, half the
     Python scan loop.
  6. ``print`` calls removed; ``return_embedding`` kept.
"""
from __future__ import annotations

import json
import math
import os

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

_DIR = os.path.dirname(os.path.abspath(__file__))
_XLSR_CFG = os.path.join(_DIR, "xlsr2_300m_cfg.json")

WINDOW = 66800  # source data_utils.py / ASVspoof2019.py: self.cut = 66800


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


# --------------------------------------------------------------------------- XLS-R
class SSLModel(nn.Module):
    def __init__(self, cfg_path: str = _XLSR_CFG):
        super().__init__()
        from fairseq.dataclass.utils import merge_with_parent
        from fairseq.models.wav2vec.wav2vec2 import Wav2Vec2Config, Wav2Vec2Model
        from omegaconf import OmegaConf

        with open(cfg_path) as f:
            stored = OmegaConf.create(json.load(f))
        cfg = merge_with_parent(Wav2Vec2Config(), stored, remove_missing=True)
        self.model = Wav2Vec2Model.build_model(cfg, task=None)
        self.out_dim = 1024

    def extract_feat(self, input_data):
        # source moves the model + calls self.model.train() on device/dtype mismatch: removed (edit 3)
        input_tmp = input_data[:, :, 0] if input_data.ndim == 3 else input_data
        return self.model(input_tmp, mask=False, features_only=True)["x"]


# --------------------------------------------------------------------------- mamba_ssm v1 subset
def selective_scan_ref(u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False):
    """mamba_ssm 1.x ops/selective_scan_interface.selective_scan_ref (real A, variable B/C).

    u, delta, z: (b, d, l); A: (d, n); B, C: (b, n, l); D, delta_bias: (d,)
    """
    dtype_in = u.dtype
    u = u.float()
    delta = delta.float()
    if delta_bias is not None:
        delta = delta + delta_bias[..., None].float()
    if delta_softplus:
        delta = F.softplus(delta)
    batch, dim, dstate = u.shape[0], A.shape[0], A.shape[1]
    B = B.float()
    C = C.float()
    x = A.new_zeros((batch, dim, dstate))
    ys = []
    deltaA = torch.exp(torch.einsum("bdl,dn->bdln", delta, A))
    deltaB_u = torch.einsum("bdl,bnl,bdl->bdln", delta, B, u)
    for i in range(u.shape[2]):
        x = deltaA[:, :, i] * x + deltaB_u[:, :, i]
        y = torch.einsum("bdn,bn->bd", x, C[:, :, i])
        ys.append(y)
    y = torch.stack(ys, dim=2)  # (b d l)
    out = y if D is None else y + u * rearrange(D, "d -> d 1")
    if z is not None:
        out = out * F.silu(z)
    return out.to(dtype=dtype_in)


def selective_scan_parallel(u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False):
    """Same recurrence as ``selective_scan_ref`` (h_t = exp(dt_t A) h_{t-1} + dt_t B_t u_t,
    y_t = C_t . h_t), evaluated as a chunked parallel scan: a Hillis-Steele prefix scan
    inside chunks of _SCAN_CHUNK steps (log2 levels of pure tensor ops), then a short
    sequential carry across chunks -- instead of an L-step Python loop. The combine
    (a1, b1) o (a2, b2) = (a1 a2, a2 b1 + b2) only multiplies decays <= 1, so it is
    numerically stable; results equal the sequential scan up to float rounding.
    """
    dtype_in = u.dtype
    u = u.float()
    delta = delta.float()
    if delta_bias is not None:
        delta = delta + delta_bias[..., None].float()
    if delta_softplus:
        delta = F.softplus(delta)
    L = u.shape[2]
    K = _SCAN_CHUNK
    nc = -(-L // K)
    # time-major layout (l b d n): shifted slices are contiguous blocks
    a = torch.exp(torch.einsum("bdl,dn->lbdn", delta, A))
    h = torch.einsum("bdl,bnl,bdl->lbdn", delta, B.float(), u)
    if nc * K != L:  # pad the tail with identity steps (a=1, h=0); truncated below
        a = torch.cat([a, a.new_ones((nc * K - L,) + a.shape[1:])])
        h = torch.cat([h, h.new_zeros((nc * K - L,) + h.shape[1:])])
    a = a.reshape((nc, K) + a.shape[1:])
    h = h.reshape((nc, K) + h.shape[1:])
    # 1) Hillis-Steele inclusive scan inside each chunk (dim 1): log2(K) levels
    k = 1
    while k < K:
        # positions < k have no earlier state in the chunk; a -> within-chunk prefix product
        h = torch.cat([h[:, :k], h[:, k:] + a[:, k:] * h[:, :-k]], dim=1)
        a = torch.cat([a[:, :k], a[:, k:] * a[:, :-k]], dim=1)
        k *= 2
    # 2) carry the state across chunks (nc short steps), then 3) add it everywhere
    carry = [h.new_zeros(h.shape[2:])]
    for c in range(nc - 1):
        carry.append(h[c, -1] + a[c, -1] * carry[-1])
    h = h + a * torch.stack(carry)[:, None]
    h = h.reshape((nc * K,) + h.shape[2:])[:L]
    y = torch.einsum("lbdn,bnl->bdl", h, C.float())
    out = y if D is None else y + u * rearrange(D, "d -> d 1")
    if z is not None:
        out = out * F.silu(z)
    return out.to(dtype=dtype_in)


_SCAN_CHUNK = 8  # 2026-10-30: fastest of 8/16/32 on RTX 4070 Ti SUPER (B=16: 127 vs 143 vs 194 ms per encoder pass)

# "parallel" (default; fast, exportable) or "ref" (the library's sequential reference loop)
SCAN_IMPL = os.environ.get("FAKE_MAMBA_SCAN", "parallel")


class Mamba(nn.Module):
    """mamba_ssm 1.x modules/mamba_simple.Mamba (forward, use_fast_path=False branch)."""

    def __init__(self, d_model, d_state=16, d_conv=4, expand=2, dt_rank="auto",
                 dt_min=0.001, dt_max=0.1, dt_init="random", dt_scale=1.0, dt_init_floor=1e-4,
                 conv_bias=True, bias=False):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.d_conv = d_conv
        self.expand = expand
        self.d_inner = int(self.expand * self.d_model)
        self.dt_rank = math.ceil(self.d_model / 16) if dt_rank == "auto" else dt_rank

        self.in_proj = nn.Linear(self.d_model, self.d_inner * 2, bias=bias)
        self.conv1d = nn.Conv1d(self.d_inner, self.d_inner, bias=conv_bias, kernel_size=d_conv,
                                groups=self.d_inner, padding=d_conv - 1)
        self.activation = "silu"
        self.act = nn.SiLU()
        self.x_proj = nn.Linear(self.d_inner, self.dt_rank + self.d_state * 2, bias=False)
        self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True)
        # (init values are irrelevant: every tensor is overwritten by the checkpoint)
        A = repeat(torch.arange(1, self.d_state + 1, dtype=torch.float32), "n -> d n", d=self.d_inner)
        self.A_log = nn.Parameter(torch.log(A).contiguous())
        self.D = nn.Parameter(torch.ones(self.d_inner))
        self.out_proj = nn.Linear(self.d_inner, self.d_model, bias=bias)

    def forward(self, hidden_states):
        batch, seqlen, dim = hidden_states.shape
        xz = rearrange(self.in_proj.weight @ rearrange(hidden_states, "b l d -> d (b l)"),
                       "d (b l) -> b d l", l=seqlen)
        if self.in_proj.bias is not None:
            xz = xz + rearrange(self.in_proj.bias.to(dtype=xz.dtype), "d -> d 1")
        A = -torch.exp(self.A_log.float())  # (d_inner, d_state)
        x, z = xz.chunk(2, dim=1)
        x = self.act(self.conv1d(x)[..., :seqlen])
        x_dbl = self.x_proj(rearrange(x, "b d l -> (b l) d"))  # (bl d)
        dt, B, C = torch.split(x_dbl, [self.dt_rank, self.d_state, self.d_state], dim=-1)
        dt = self.dt_proj.weight @ dt.t()
        dt = rearrange(dt, "d (b l) -> b d l", l=seqlen)
        B = rearrange(B, "(b l) dstate -> b dstate l", l=seqlen).contiguous()
        C = rearrange(C, "(b l) dstate -> b dstate l", l=seqlen).contiguous()
        scan = selective_scan_parallel if SCAN_IMPL == "parallel" else selective_scan_ref
        y = scan(x, dt, A, B, C, self.D.float(), z=z,
                               delta_bias=self.dt_proj.bias.float(), delta_softplus=True)
        y = rearrange(y, "b d l -> b l d")
        return self.out_proj(y)


# --------------------------------------------------------------------------- PN-BiMamba
class PN_BiMambas_Encoder(nn.Module):
    def __init__(self, d_model, n_state):
        super().__init__()
        self.d_model = d_model
        self.mamba = Mamba(d_model, n_state)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.GELU(),
            nn.Linear(d_model * 4, d_model),
        )

    def forward(self, x):
        residual = x
        x_norm = self.norm1(x)
        b = x_norm.shape[0]
        # forward + backward (flipped) through the shared Mamba in one call (edit 5)
        both = self.mamba(torch.cat([x_norm, torch.flip(x_norm, dims=[1])], dim=0))
        mamba_out_forward = both[:b]
        mamba_out_backward = torch.flip(both[b:], dims=[1])
        mamba_out = mamba_out_forward + mamba_out_backward
        mamba_out = self.norm2(mamba_out)
        ff_out = self.feed_forward(mamba_out)
        return ff_out + residual


class PN_BiMambas(nn.Module):
    def __init__(self, dim, *, depth):
        super().__init__()
        self.dim = dim
        self.layers = nn.ModuleList([PN_BiMambas_Encoder(d_model=dim, n_state=16) for _ in range(depth)])

    def forward(self, x):
        for block in self.layers:
            x = block(x)
        return x


class FakeMambaNet(nn.Module):
    """model.py:Fake_Mamba (args.emb_size = 144). forward(x[B, T]) -> logits[B, 2]."""

    def __init__(self, emb_size: int = 144, depth: int = 12):
        super().__init__()
        self.ssl_model = SSLModel()
        self.LL = nn.Linear(1024, emb_size)
        self.first_bn = nn.BatchNorm2d(num_features=1)
        self.selu = nn.SELU(inplace=True)
        self.encoder = PN_BiMambas(dim=144, depth=depth)
        self.attention_pool = nn.Linear(144, 1)
        self.fc5 = nn.Linear(144, 2)

    def forward(self, x, return_embedding=False):
        x_ssl_feat = self.ssl_model.extract_feat(x.squeeze(-1))  # (B, T, 1024)
        x = self.LL(x_ssl_feat)  # (B, T, 144)
        x = x.unsqueeze(dim=1)
        x = self.first_bn(x)
        x = self.selu(x)
        x = x.squeeze(dim=1)
        x = self.encoder(x)
        attention_weights = F.softmax(self.attention_pool(x), dim=1)  # (B, T, 1)
        x_pooled = torch.matmul(attention_weights.transpose(-1, -2), x).squeeze(-2)  # (B, 144)
        if return_embedding:
            return self.fc5(x_pooled), x_pooled
        return self.fc5(x_pooled)


def load_state(path: str) -> dict:
    sd = torch.load(path, map_location="cpu")
    return sd.get("state_dict", sd) if isinstance(sd, dict) and "state_dict" in sd else sd