"""Ported from github.com/ProSDD/codes @ bdb80440cd946368e77f4b51f6ddf26c7443dc33 (no LICENSE file):
model_stage2realfake.py (ProSDDStage2), main_eval.py (inference_forward), data_utils_eval.py (pad).

Inference-only subset of ProSDDStage2. Edits for the Arena wrapper:
  1. `ssl` is built from the vendored facebook/wav2vec2-xls-r-300m config (xlsr300m_config.json, HF @
     1a640f32ac3e39899438a2931f9924c02f080a54) instead of Wav2Vec2Model.from_pretrained(...): no Hub
     download; the checkpoint overwrites every backbone tensor (strict load, see the wrapper).
  2. Training-only members are not built: mask_embed, pros_ln, final_proj (1024->448 Stage-1 head) and the
     masked/contrastive pass. classifier_pool is fixed to "mean" (the default; the released checkpoints have
     no attn/attn_q tensors).
  3. forward(wav) = main_eval.py:inference_forward verbatim (CNN feature extractor -> feature projection ->
     latent truncated / zero-padded to T_target=200 frames -> encoder (clean pass, no attention mask) ->
     mean over the 200 frames -> cls_head), returning the (B, 2) logits.
  4. pad_np = numpy port of data_utils_eval.py:pad (64000 samples, centre crop / centred zero-pad).
"""
import os

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

_DIR = os.path.dirname(os.path.abspath(__file__))
XLSR_CONFIG = os.path.join(_DIR, "xlsr300m_config.json")
TARGET_SAMPLES = 64000  # data_utils_eval.py: 4.0 s
T_TARGET = 200          # main_eval.py: ProSDDStage2(T_target=200); 64000 samples give 199 latent frames


def pad_np(wav, max_len=TARGET_SAMPLES):
    """data_utils_eval.py:pad -- longer: centre crop s=(n-max_len)//2; shorter: zero-pad left=total//2, right=rest."""
    wav = np.asarray(wav, dtype=np.float32).reshape(-1)
    n = wav.shape[0]
    if n >= max_len:
        s = (n - max_len) // 2
        return np.ascontiguousarray(wav[s:s + max_len])
    left = (max_len - n) // 2
    out = np.zeros(max_len, dtype=np.float32)
    out[left:left + n] = wav
    return out


class ProSDDNet(nn.Module):
    def __init__(self, num_classes=2, T_target=T_TARGET):
        super().__init__()
        self.T_target = int(T_target)
        self.ssl = Wav2Vec2Model(Wav2Vec2Config.from_json_file(XLSR_CONFIG))  # edit 1
        self.hidden_dim = self.ssl.config.hidden_size  # 1024
        self.cls_head = nn.Sequential(
            nn.Linear(self.hidden_dim, 512),
            nn.ReLU(inplace=True),
            nn.Dropout(0.1),
            nn.Linear(512, num_classes),
        )

    def forward(self, wav):  # (B, 64000) 16 kHz -> logits (B, 2); class 1 = bona fide
        z = self.ssl.feature_extractor(wav)
        z = z.transpose(1, 2)
        z = self.ssl.feature_projection(z)
        if isinstance(z, (tuple, list)):
            z = z[0]
        B, T, H = z.shape
        Tt = self.T_target
        if T > Tt:
            z = z[:, :Tt, :]
        elif T < Tt:
            z = torch.cat([z, z.new_zeros(B, Tt - T, H)], dim=1)
        out = self.ssl.encoder(z, attention_mask=None, output_hidden_states=False, return_dict=True)
        pooled = out.last_hidden_state.mean(dim=1)
        return self.cls_head(pooled)