"""Vendored from github.com/adipiz99/RawNetLite @ fcb0fe8d0f40146d110e322c00a0723ba27d5520 (MIT, (c) 2025 Andrea Di Pierno):
RawNetLite.py (ResBlock, RawNetLite) and audio_preprocessor.py (preprocess_audio).

Edits for the Arena wrapper:
  1. RawNetLite.forward_logit = forward without the final torch.sigmoid (returns the fc2 pre-activation, the
     spoof logit); forward() is kept verbatim (sigmoid(spoof logit) = P(fake)). Docstrings dropped.
  2. prep_np = numpy port of preprocess_audio for an already-mono 16 kHz clip: divide by max |x| over the WHOLE clip,
     then keep the first 48000 samples (3 s) / right zero-pad. No resampling (the Arena delivers 16 kHz =
     the authors' target_sr). A silent/empty clip (max |x| = 0, where the source would give NaN) -> zeros.
"""
import numpy as np
import torch
import torch.nn as nn

TARGET_SR = 16000
TARGET_SAMPLES = 48000  # audio_preprocessor.py: target_sec=3.0 at 16 kHz (shorter than the Arena 4 s cap)


def prep_np(wav, max_len=TARGET_SAMPLES):
    wav = np.asarray(wav, dtype=np.float32).reshape(-1)
    peak = np.abs(wav).max() if wav.size else 0.0
    if peak > 0:
        wav = wav / peak
    out = np.zeros(max_len, dtype=np.float32)
    n = min(wav.shape[0], max_len)
    out[:n] = wav[:n]
    return out


class ResBlock(nn.Module):
    def __init__(self, channels):
        super(ResBlock, self).__init__()
        self.conv1 = nn.Conv1d(channels, channels, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm1d(channels)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv1d(channels, channels, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm1d(channels)

    def forward(self, x):
        residual = x
        x = self.relu(self.bn1(self.conv1(x)))
        x = self.bn2(self.conv2(x))
        return self.relu(x + residual)


class RawNetLite(nn.Module):
    def __init__(self):
        super(RawNetLite, self).__init__()
        self.conv_pre = nn.Conv1d(1, 64, kernel_size=3, stride=1, padding=1)
        self.bn_pre = nn.BatchNorm1d(64)
        self.relu = nn.ReLU()

        self.resblock1 = ResBlock(64)
        self.resblock2 = ResBlock(64)
        self.resblock3 = ResBlock(64)

        self.pool = nn.AdaptiveAvgPool1d(64)  # Sequence reduction for GRU

        self.gru = nn.GRU(input_size=64, hidden_size=128, num_layers=1,
                          batch_first=True, bidirectional=True)

        self.fc1 = nn.Linear(128 * 2, 64)
        self.fc2 = nn.Linear(64, 1)

    def forward_logit(self, x):  # edit 1
        # x: [B, 1, T]
        x = self.relu(self.bn_pre(self.conv_pre(x)))     # [B, 64, T]
        x = self.resblock1(x)
        x = self.resblock2(x)
        x = self.resblock3(x)
        x = self.pool(x)                                 # [B, 64, 64]

        x = x.transpose(1, 2)                            # [B, 64, 64] -> [B, seq, feat]
        output, _ = self.gru(x)                          # [B, 64, 256]
        x = output[:, -1, :]                             # Last step -> [B, 256]

        x = self.fc1(x)                                  # [B, 64]
        x = self.fc2(x)                                  # [B, 1]
        return x

    def forward(self, x):
        return torch.sigmoid(self.forward_logit(x))