"""DeepFense StandardDetector (wav2vec2 frontend + MLP backend + CrossEntropy head), inference only.

Ported from github.com/Yaselley/deepfense-framework @ 3f2fe1146c48ea4bacc480dd96d806b6c72c09cf (Apache-2.0):
  deepfense/models/detector.py            ModularDetector (frontend -> backend -> losses[0])
  deepfense/models/frontends/wav2vec2.py  Wav2VecWrapper, source=fairseq: model(x, mask=False, features_only=True)["x"]
  deepfense/models/backends/mlp.py        MLP: projection [128] -> Linear(1024,128), LayerNorm(128), ReLU; mean pool
  deepfense/models/losses/cross_entropy.py fc = Linear(128, 2)
  deepfense/data/transforms/transforms.py pad(max_len=64000, pad_type=repeat, random_pad=False at eval)
with the hyper-parameters of the checkpoint's config.yaml (HF DeepFense/CodecFake_Wav2Vec2_MLP_NoAug_Seed42
@ c08bd2125c83e6c8fdc64566bfbfda18dff4dfb4).

Edits for the Arena wrapper:
  1. The XLS-R 300M module is the installed fairseq (0.12.2) Wav2Vec2Model built from the vendored model cfg of
     xlsr2_300m.pt (xlsr2_300m_model_cfg.json, merged onto Wav2Vec2Config defaults as in W2V2-AASIST/_net.py)
     instead of loading the 3.6 GB pre-trained .pt. The checkpoint's model_state was saved from this fairseq
     module (it carries quantizer / project_q / final_proj), so all 435 tensors map 1:1 (strict load). Those
     pre-training heads are unused on the features_only path.
  2. Module names reproduce the checkpoint keys: frontend.model.*, backend.projection_block.{0,1}.*, losses.0.fc.*.
  3. forward returns the (B, 2) logits only (no embeddings/probs dict); the mean pool is written inline
     (pool.py "mean" == x.mean over time on [B, D, T]).
"""
import json
import os

import numpy as np
import torch
import torch.nn as nn

_DIR = os.path.dirname(os.path.abspath(__file__))
XLSR_CFG = os.path.join(_DIR, "xlsr2_300m_model_cfg.json")
WINDOW = 64000  # config.yaml data.test.base_transform: pad max_len 64000, pad_type repeat


def pad_np(wav, n=WINDOW):
    """transforms.py pad_combined at eval: first n samples, else np.tile(x, ceil(n/len))[:n]. No normalisation."""
    x = np.asarray(wav, dtype=np.float32).reshape(-1)
    if x.shape[0] == 0:  # (not in source) empty clip -> silence
        return np.zeros(n, dtype=np.float32)
    if x.shape[0] > n:
        return np.ascontiguousarray(x[:n])
    return np.ascontiguousarray(np.tile(x, int(np.ceil(n / x.shape[0])))[:n])


def build_xlsr():
    from fairseq.dataclass.utils import merge_with_parent
    from fairseq.models.wav2vec.wav2vec2 import Wav2Vec2Config, Wav2Vec2Model
    from omegaconf import OmegaConf

    cfg = OmegaConf.create(json.load(open(XLSR_CFG)))
    return Wav2Vec2Model.build_model(merge_with_parent(Wav2Vec2Config(), cfg, remove_missing=True), task=None)


class _Frontend(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = build_xlsr()

    def forward(self, x):
        return self.model(x, mask=False, features_only=True)["x"]  # (B, T', 1024), last layer


class _MLPBackend(nn.Module):
    def __init__(self, in_dim=1024, dim=128):
        super().__init__()
        self.projection_block = nn.Sequential(nn.Linear(in_dim, dim), nn.LayerNorm(dim), nn.ReLU(inplace=True))

    def forward(self, x):
        return self.projection_block(x).transpose(1, 2).mean(dim=-1)  # (B, 128)


class _CEHead(nn.Module):
    def __init__(self, dim=128, n_classes=2):
        super().__init__()
        self.fc = nn.Linear(dim, n_classes)

    def forward(self, e):
        return self.fc(e)


class DeepFenseDetector(nn.Module):
    def __init__(self):
        super().__init__()
        self.frontend = _Frontend()
        self.backend = _MLPBackend()
        self.losses = nn.ModuleList([_CEHead()])

    def forward(self, waveform):  # (B, 64000) -> logits (B, 2); label_map bonafide=1, spoof=0
        return self.losses[0](self.backend(self.frontend(waveform)))