from __future__ import annotations

import os

import numpy as np
import torch
import torch.nn.functional as F
from speech_spoof_bench.model import AntiSpoofingModel

from _net import WINDOW, SSLMambaNet, load_state, pad_repeat

_DIR = os.path.dirname(os.path.abspath(__file__))
# HF hoanmyTran/ssl_spoofing_mamba @ 4e72dfc interspeech2025.pt -> safetensors (loss_fn.weight dropped)
_CKPT = os.path.join(_DIR, "model.safetensors")


class SSLMamba(AntiSpoofingModel):
    name = "SSL-Mamba"
    expected_sample_rate = 16000
    batch_size = 16  # sweep 2026-11-01 (pure torch, GPUs shared): RTX 4070 Ti SUPER ~57 utt/s, plateau bs 4-32

    def load(self) -> None:
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        net = SSLMambaNet()
        net.load_state_dict(load_state(_CKPT), strict=True)
        self.net = net.to(self.device).eval()
        if self.device == "cuda" and "5060" in torch.cuda.get_device_name():
            self.batch_size = 4  # sweep 2026-11-01: RTX 5060 Ti ~28.6 utt/s at bs 2-4, ~22 at bs >= 8

    @torch.no_grad()
    def score_batch(self, audios, srs):
        x = np.stack([pad_repeat(a, WINDOW) for a in audios])
        out = self.net(torch.from_numpy(x).to(self.device)).float().reshape(-1, 2)
        ls = F.log_softmax(out, dim=-1)
        # source _produce_evaluation_file: llr = log_softmax[:,1] - log_softmax[:,0]; class 1 = bona fide
        return (ls[:, 1] - ls[:, 0]).cpu().tolist()

    def unload(self) -> None:
        self.net = None
        if torch.cuda.is_available():
            torch.cuda.empty_cache()