"""LaST-Net (Length-Aware Layer and Scale-Adaptive Temporal Network), inference only.

Vendored from github.com/JunXue-tech/LaST-Net @ 3838f0c7108313f771f6c1e208ac8f1e3bde60b0 (MIT, Copyright (c) 2026
LaST-Net authors; see LICENSE.source):
  model.py         duration_features, LengthCalibratedLayerSelector, TemporalResidualBlock,
                   LengthSelectiveKernelBlock, LengthPriorResidualExchange, AttentiveStatisticsPooling, LaSTNet
  load_dataset.py  crop_or_repeat (prefix crop / np.tile repeat), `available` seconds = min(len, samples) / 16000
  infer.py/eval.py score = log_softmax(logits)[:, 1] ("bonafide_log_score"; label bonafide = 1)

Edits for the Arena wrapper (math unchanged):
  1. SSLModel builds the installed fairseq 0.12.2 Wav2Vec2Model 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
     load_model_ensemble_and_task([xlsr2_300m.pt]); the fine-tuned checkpoint overwrites every tensor (strict load).
     extract_feat is verbatim: model(wav, mask=False, features_only=True, layer=23)["layer_results"].
  2. LaSTNet.__init__ takes no ssl_path; forward returns the (B, 2) logits only (return_aux dropped).
  3. crop_np = crop_or_repeat (plus an empty-clip -> silence guard, which the source rejects with ValueError).
  4. nn.init.zeros_ calls dropped (every tensor is overwritten by the strict load); code reflowed.
Everything else is copied verbatim (the training-only layer-dropout branch is inert in eval()).
"""
import json
import math
import os

import numpy as np
import torch
from torch import nn
from torch.nn import functional as F

_DIR = os.path.dirname(os.path.abspath(__file__))
XLSR_CFG = os.path.join(_DIR, "xlsr2_300m_model_cfg.json")
SAMPLE_RATE = 16000
SECONDS = 4                       # infer.py --seconds 4 (Arena 4 s cap)
WINDOW = SECONDS * SAMPLE_RATE    # 64000: infer.py/eval.py use seconds * 16000


def crop_np(wav, samples=WINDOW):
    """load_dataset.crop_or_repeat + the `available` seconds that eval.py/infer.py pass to the model."""
    x = np.asarray(wav, dtype=np.float32).reshape(-1)
    if x.shape[0] == 0:  # (not in source: it raises) empty clip -> silence, available = 0 (clamped to 1 ms)
        return np.zeros(samples, dtype=np.float32), 0.0
    available = min(len(x), samples) / SAMPLE_RATE
    if len(x) < samples:
        x = np.tile(x, math.ceil(samples / len(x)))
    return np.ascontiguousarray(x[:samples], dtype=np.float32), available


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 SSLModel(nn.Module):
    """XLS-R frontend returning all 24 transformer-layer representations."""

    def __init__(self):
        super().__init__()
        self.model = build_xlsr()  # edit 1

    def extract_feat(self, wav):
        output = self.model(wav, mask=False, features_only=True, layer=23)
        return output["layer_results"]


def duration_features(wav, available_seconds=None):
    requested = wav.new_full((len(wav),), wav.shape[-1] / 16000.0)
    available = requested if available_seconds is None else available_seconds.to(wav)
    available = available.clamp(0.001, 6.0)
    return torch.stack(
        (
            requested.clamp_max(6.0) / 6.0,
            torch.log1p(requested.clamp_max(6.0)) / torch.log(requested.new_tensor(7.0)),
            available / requested.clamp_min(0.001),
        ),
        dim=1,
    )


class LengthCalibratedLayerSelector(nn.Module):
    """Fuse the 24 XLS-R layers using length and content-conditioned weights."""

    def __init__(self, dimension=1024, temperature=0.8, uniform_floor=0.05, layer_dropout=0.10):
        super().__init__()
        self.temperature = temperature
        self.uniform_floor = uniform_floor
        self.layer_dropout = layer_dropout
        self.content_scorer = nn.Sequential(
            nn.LayerNorm(2 * dimension), nn.Linear(2 * dimension, 128), nn.GELU(), nn.Linear(128, 1),
        )
        self.length_scorer = nn.Sequential(nn.Linear(3, 64), nn.GELU(), nn.Linear(64, 24))
        self.layer_bias = nn.Parameter(torch.zeros(24))

    def forward(self, layers, lengths):
        if len(layers) != 24:
            raise RuntimeError("LaST-Net expects all 24 XLS-R layers")
        sequences = [layer[0].transpose(0, 1) for layer in layers]
        descriptors = []
        for sequence in sequences:
            mean = sequence.mean(dim=1)
            std = (sequence - mean.unsqueeze(1)).square().mean(dim=1).sqrt()
            descriptors.append(torch.cat((mean, std), dim=-1))
        descriptors = torch.stack(descriptors, dim=1)

        content_scores = self.content_scorer(descriptors.float()).squeeze(-1)
        length_scores = self.length_scorer(lengths.float())
        scores = (self.layer_bias.unsqueeze(0) + content_scores + length_scores) / self.temperature

        if self.training and self.layer_dropout > 0.0:
            if len(scores) % 2 == 0:
                mask = torch.rand(len(scores) // 2, 24, device=scores.device) < self.layer_dropout
                mask = mask.repeat(2, 1)
            else:
                mask = torch.rand_like(scores) < self.layer_dropout
            scores = scores.masked_fill(mask, -1e4)

        conditional = F.softmax(scores, dim=1)
        weights = (1.0 - self.uniform_floor) * conditional + self.uniform_floor / 24.0
        fused = sum(sequence * weights[:, index, None, None].to(sequence) for index, sequence in enumerate(sequences))
        return fused, weights


class TemporalResidualBlock(nn.Module):
    def __init__(self, dimension, dilation):
        super().__init__()
        self.norm = nn.LayerNorm(dimension)
        self.depthwise = nn.Conv1d(dimension, dimension, 5, padding=2 * dilation, dilation=dilation,
                                   groups=dimension, bias=False)
        self.pointwise = nn.Conv1d(dimension, 2 * dimension, 1)
        self.output = nn.Linear(dimension, dimension)
        self.scale = nn.Parameter(torch.full((dimension,), 0.1))

    def forward(self, x):
        y = self.depthwise(self.norm(x).transpose(1, 2))
        value, gate = self.pointwise(y).chunk(2, dim=1)
        update = (value * torch.sigmoid(gate)).transpose(1, 2)
        return x + self.scale * self.output(update)


class LengthSelectiveKernelBlock(nn.Module):
    """Length-guided local temporal filtering at three receptive fields."""

    def __init__(self, dimension):
        super().__init__()
        self.norm = nn.LayerNorm(dimension)
        self.convs = nn.ModuleList([
            nn.Conv1d(dimension, dimension, 5, padding=2 * dilation, dilation=dilation, groups=dimension, bias=False)
            for dilation in (1, 2, 4)
        ])
        self.content_gate = nn.Linear(dimension, 2)
        self.pointwise = nn.Conv1d(dimension, 2 * dimension, 1)
        self.output = nn.Linear(dimension, dimension)
        self.scale = nn.Parameter(torch.full((dimension,), 0.1))

    def forward(self, x, lengths):
        normalized = self.norm(x)
        channels = normalized.transpose(1, 2)
        features = [conv(channels) for conv in self.convs]
        duration = lengths[:, :1].float()
        residual = 0.1 * torch.tanh(self.content_gate(normalized.mean(dim=1).float()))
        gate2 = torch.sigmoid((duration - 0.35) / 0.15 + residual[:, :1]).to(x)
        gate4 = torch.sigmoid((duration - 0.65) / 0.15 + residual[:, 1:]).to(x)
        mixed = (
            features[0] + gate2.unsqueeze(-1) * features[1] + gate4.unsqueeze(-1) * features[2]
        ) / (1.0 + gate2 + gate4).unsqueeze(-1)
        value, gate = self.pointwise(mixed).chunk(2, dim=1)
        update = (value * torch.sigmoid(gate)).transpose(1, 2)
        return x + self.scale * self.output(update)


class LengthPriorResidualExchange(nn.Module):
    """Process three temporal scales and exchange length-relevant evidence."""

    def __init__(self, dimension, selective_local=False):
        super().__init__()
        self.selective_local = selective_local
        self.blocks = nn.ModuleList([
            LengthSelectiveKernelBlock(dimension) if selective_local else TemporalResidualBlock(dimension, 1),
            TemporalResidualBlock(dimension, 2),
            TemporalResidualBlock(dimension, 4),
        ])
        self.edge_bias = nn.Parameter(torch.zeros(6))
        self.content_gate = nn.Linear(dimension, 6)
        self.fusion_scale = nn.Parameter(torch.full((3, dimension), 0.1))

    @staticmethod
    def resize(x, size):
        if x.shape[1] == size:
            return x
        return F.interpolate(x.transpose(1, 2), size=size, mode="linear", align_corners=False).transpose(1, 2)

    def forward(self, scales, lengths):
        processed = []
        for index, (block, scale) in enumerate(zip(self.blocks, scales)):
            if index == 0 and self.selective_local:
                processed.append(block(scale, lengths))
            else:
                processed.append(block(scale))

        duration = lengths[:, 0].clamp(0.0, 1.0)
        source_prior = torch.stack((1.0 - duration, 4.0 * duration * (1.0 - duration), duration), dim=1)
        content = sum(x.mean(dim=1) for x in processed) / 3.0
        content_residual = 0.1 * torch.tanh(self.content_gate(content.float()))

        outputs = []
        all_weights = []
        edge_index = 0
        for target in range(3):
            sources = [source for source in range(3) if source != target]
            logits = []
            for local_index, source in enumerate(sources):
                index = edge_index + local_index
                logits.append(self.edge_bias[index] + 2.0 * (source_prior[:, source] - 0.5)
                              + content_residual[:, index])
            positive = F.softplus(torch.stack(logits, dim=1)).to(processed[target])
            weights = positive / positive.sum(dim=1, keepdim=True).clamp_min(1e-6)
            incoming = sum(
                weights[:, index, None, None] * self.resize(processed[source], processed[target].shape[1])
                for index, source in enumerate(sources)
            )
            outputs.append(processed[target] + self.fusion_scale[target] * incoming)
            all_weights.append(weights)
            edge_index += 2
        return outputs, torch.cat(all_weights, dim=1)


class AttentiveStatisticsPooling(nn.Module):
    def __init__(self, dimension=192):
        super().__init__()
        self.attention = nn.Sequential(nn.Linear(dimension, 128), nn.Tanh(), nn.Linear(128, 1))

    def forward(self, x):
        weights = F.softmax(self.attention(x).float(), dim=1).to(x)
        mean = (weights * x).sum(dim=1)
        second = (weights * x.square()).sum(dim=1)
        std = (second - mean.square()).clamp_min(1e-5).sqrt()
        return torch.cat((mean, std), dim=-1)


class LaSTNet(nn.Module):
    architecture = "LaST-Net"

    # checkpoint args: router_temperature 0.8, uniform_floor 0.05, layer_dropout 0.1 (eval.py load_model reads these)
    def __init__(self, temperature=0.8, uniform_floor=0.05, layer_dropout=0.10, dimension=192):
        super().__init__()
        self.ssl_model = SSLModel()
        self.router = LengthCalibratedLayerSelector(temperature=temperature, uniform_floor=uniform_floor,
                                                    layer_dropout=layer_dropout)
        self.input_projection = nn.Sequential(nn.Linear(1024, dimension), nn.LayerNorm(dimension))
        self.exchanges = nn.ModuleList([
            LengthPriorResidualExchange(dimension, selective_local=True),
            LengthPriorResidualExchange(dimension, selective_local=False),
        ])
        self.pools = nn.ModuleList([AttentiveStatisticsPooling(dimension) for _ in range(3)])
        self.classifier = nn.Sequential(
            nn.LayerNorm(6 * dimension), nn.Linear(6 * dimension, 256), nn.GELU(), nn.Dropout(0.2), nn.Linear(256, 2),
        )

    def forward(self, wav, available_seconds=None):  # (B, T) 16 kHz, (B,) seconds -> logits (B, 2); class 1 = bona fide
        layers = self.ssl_model.extract_feat(wav)
        lengths = duration_features(wav, available_seconds)
        fused, _ = self.router(layers, lengths)
        x0 = self.input_projection(fused)
        scales = [
            x0,
            F.avg_pool1d(x0.transpose(1, 2), 2, 2).transpose(1, 2),
            F.avg_pool1d(x0.transpose(1, 2), 4, 4).transpose(1, 2),
        ]
        for exchange in self.exchanges:
            scales, _ = exchange(scales, lengths)
        embedding = torch.cat([pool(scale) for pool, scale in zip(self.pools, scales)], dim=-1)
        return self.classifier(embedding)