"""Telephony Audio Degradation Pipeline for ASVspoof 5.

Simulates realistic voice call channel artifacts:
1. PSTN Narrowband Bandpass Filtering (300 Hz - 3400 Hz).
2. 8 kHz Downsampling (Nyquist = 4000 Hz).
3. G.711 Mu-law 8-bit Companding and Quantization.
4. VoIP / RTP Packet Loss Simulation (frame erasure / burst drops).
"""

from typing import Optional, Tuple
import torch
import torch.nn as nn
import torchaudio.transforms as T
import torchaudio.functional as F


class TelephonyBandpassFilter(nn.Module):
    """Simulates the 300 Hz - 3400 Hz acoustic passband of standard PSTN telephone networks."""

    def __init__(self, sample_rate: int = 16000, low_freq: float = 300.0, high_freq: float = 3400.0):
        super().__init__()
        self.sample_rate = sample_rate
        self.low_freq = low_freq
        self.high_freq = high_freq

    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
        """Apply highpass at 300 Hz and lowpass at 3400 Hz.
        
        waveform: Tensor of shape (..., num_samples)
        """
        # Highpass to remove low-frequency rumble / DC offset (< 300 Hz)
        filtered = F.highpass_biquad(waveform, sample_rate=self.sample_rate, cutoff_freq=self.low_freq)
        # Lowpass to limit to narrowband telephone bandwidth (< 3400 Hz)
        filtered = F.lowpass_biquad(filtered, sample_rate=self.sample_rate, cutoff_freq=self.high_freq)
        return filtered


class PacketLossSimulator(nn.Module):
    """Simulates VoIP / RTP burst packet loss with frame erasure and packet loss concealment."""

    def __init__(
        self,
        sample_rate: int = 8000,
        packet_duration_ms: float = 20.0,
        loss_rate: float = 0.05,
        max_burst_packets: int = 3,
        concealment: str = "zero"  # 'zero', 'noise', or 'attenuate'
    ):
        super().__init__()
        self.sample_rate = sample_rate
        self.packet_size = int(sample_rate * (packet_duration_ms / 1000.0))
        self.loss_rate = loss_rate
        self.max_burst_packets = max_burst_packets
        self.concealment = concealment

    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
        """Apply packet drops on 20ms frames.
        
        waveform: Tensor of shape (..., num_samples)
        """
        if self.loss_rate <= 0.0 or not self.training:
            return waveform

        out = waveform.clone()
        num_samples = out.shape[-1]
        num_packets = num_samples // self.packet_size

        if num_packets == 0:
            return out

        i = 0
        while i < num_packets:
            if torch.rand(1).item() < self.loss_rate:
                burst_len = torch.randint(1, self.max_burst_packets + 1, (1,)).item()
                burst_len = min(burst_len, num_packets - i)
                start_sample = i * self.packet_size
                end_sample = min((i + burst_len) * self.packet_size, num_samples)

                if self.concealment == "zero":
                    out[..., start_sample:end_sample] = 0.0
                elif self.concealment == "noise":
                    comfort_noise = torch.randn_like(out[..., start_sample:end_sample]) * 0.001
                    out[..., start_sample:end_sample] = comfort_noise
                elif self.concealment == "attenuate":
                    out[..., start_sample:end_sample] *= 0.1

                i += burst_len
            else:
                i += 1

        return out


class TelephonyDegradation(nn.Module):
    """Complete telephony degradation pipeline.
    
    Transforms 16 kHz input speech into realistic 8 kHz narrowband telephony speech:
    1. Bandpass filter: 300 - 3400 Hz.
    2. Resampling: 16 kHz -> 8 kHz.
    3. G.711 Mu-law companding & 8-bit quantization.
    4. VoIP packet loss simulation (during training).
    5. Amplitude normalization.
    """

    def __init__(
        self,
        orig_sample_rate: int = 16000,
        target_sample_rate: int = 8000,
        apply_bandpass: bool = True,
        apply_mulaw: bool = True,
        packet_loss_rate: float = 0.05,
        packet_duration_ms: float = 20.0
    ):
        super().__init__()
        self.orig_sample_rate = orig_sample_rate
        self.target_sample_rate = target_sample_rate
        self.apply_bandpass = apply_bandpass
        self.apply_mulaw = apply_mulaw

        # 1. Bandpass (300 - 3400 Hz at 16 kHz)
        self.bandpass = TelephonyBandpassFilter(orig_sample_rate, 300.0, 3400.0)

        # 2. Resampler (16 kHz -> 8 kHz)
        self.resampler = T.Resample(orig_freq=orig_sample_rate, new_freq=target_sample_rate)

        # 3. G.711 Mu-law (8-bit, 256 quantization channels)
        self.mulaw_encoder = T.MuLawEncoding(quantization_channels=256)
        self.mulaw_decoder = T.MuLawDecoding(quantization_channels=256)

        # 4. Packet Loss Simulator (at 8 kHz)
        self.packet_loss = PacketLossSimulator(
            sample_rate=target_sample_rate,
            packet_duration_ms=packet_duration_ms,
            loss_rate=packet_loss_rate
        )

    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
        """Apply telephony degradation pipeline.
        
        Args:
            waveform: Tensor of shape (channels, num_samples) or (batch, channels, num_samples)
                      Sampled at orig_sample_rate (16 kHz).
        Returns:
            degraded: Tensor at target_sample_rate (8 kHz) in [-1.0, 1.0].
        """
        # Ensure float32
        x = waveform.float()

        # Step 1: Bandpass filter (PSTN standard passband 300-3400 Hz)
        if self.apply_bandpass:
            x = self.bandpass(x)

        # Step 2: Resample to 8 kHz
        x = self.resampler(x)

        # Step 3: G.711 Mu-law companding & 8-bit non-linear quantization
        if self.apply_mulaw:
            # Clamp to [-1.0, 1.0] before mu-law companding
            x = torch.clamp(x, -1.0, 1.0)
            encoded = self.mulaw_encoder(x)
            x = self.mulaw_decoder(encoded)

        # Step 4: VoIP packet loss simulation
        x = self.packet_loss(x)

        # Step 5: Peak normalize to avoid clipping while maintaining dynamics
        max_val = torch.max(torch.abs(x), dim=-1, keepdim=True)[0]
        max_val = torch.clamp(max_val, min=1e-6)
        # Avoid amplifying pure silence
        scale = torch.where(max_val > 0.01, 0.95 / max_val, torch.ones_like(max_val))
        x = x * scale

        return x


class TandemTelephonyDegradation(nn.Module):
    """Dual-Codec (Tandem) Telephony Degradation Pipeline.
    
    Simulates real-world phone calls traversing multiple heterogeneous networks:
    1. Bandpass filter: 300 - 3400 Hz.
    2. Resampling: 16 kHz -> 8 kHz.
    3. Step 1 (Cellular Carrier): GSM-FR 6.10 (13 kbps, 8 kHz RPE-LTP).
    4. Step 2 (VoIP / Web Gateway): G.711 µ-law or Opus compression.
    5. VoIP packet loss simulation (burst drops on 20ms frames).
    6. Peak normalization.
    """

    def __init__(
        self,
        orig_sample_rate: int = 16000,
        target_sample_rate: int = 8000,
        cellular_codec: str = "GSM610",  # GSM-FR 6.10
        voip_codec: str = "ULAW",        # G.711 µ-law or 'OPUS'
        packet_loss_rate: float = 0.05,
        packet_duration_ms: float = 20.0
    ):
        super().__init__()
        self.orig_sample_rate = orig_sample_rate
        self.target_sample_rate = target_sample_rate
        self.cellular_codec = cellular_codec
        self.voip_codec = voip_codec.upper()

        self.bandpass = TelephonyBandpassFilter(orig_sample_rate, 300.0, 3400.0)
        self.resampler = T.Resample(orig_freq=orig_sample_rate, new_freq=target_sample_rate)
        self.packet_loss = PacketLossSimulator(
            sample_rate=target_sample_rate,
            packet_duration_ms=packet_duration_ms,
            loss_rate=packet_loss_rate
        )

    def _compress_gsm_single(self, audio_np: "np.ndarray") -> "np.ndarray":
        """In-memory GSM 6.10 Full Rate compression/decompression via libsndfile."""
        import io
        import soundfile as sf
        bio = io.BytesIO()
        sf.write(bio, audio_np, self.target_sample_rate, format="WAV", subtype="GSM610")
        bio.seek(0)
        decoded, _ = sf.read(bio, dtype="float32")
        return decoded

    def _compress_voip_single(self, audio_np: "np.ndarray") -> "np.ndarray":
        """In-memory VoIP compression (G.711 ULAW or OPUS) via libsndfile."""
        import io
        import soundfile as sf
        bio = io.BytesIO()
        if self.voip_codec == "OPUS":
            sf.write(bio, audio_np, self.target_sample_rate, format="OGG", subtype="OPUS")
        else:
            sf.write(bio, audio_np, self.target_sample_rate, format="WAV", subtype="ULAW")
        bio.seek(0)
        decoded, _ = sf.read(bio, dtype="float32")
        return decoded

    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
        """Apply sequential tandem compression (Cellular GSM-FR -> VoIP ULAW/Opus).
        
        Args:
            waveform: Tensor of shape (..., num_samples) at 16 kHz.
        Returns:
            degraded: Tensor of shape (..., target_samples) at 8 kHz.
        """
        import numpy as np

        x = waveform.float()

        # Step 1: Bandpass filter (300 - 3400 Hz)
        x = self.bandpass(x)

        # Step 2: Downsample to 8 kHz
        x = self.resampler(x)

        # Flatten batch dimensions for in-memory codec passes
        orig_shape = x.shape
        x_flat = x.reshape(-1, orig_shape[-1])  # (B, N)
        target_len = orig_shape[-1]

        processed_batch = []
        for i in range(x_flat.shape[0]):
            sig = x_flat[i].detach().cpu().numpy()
            # Normalize to [-0.95, 0.95] to prevent codec clipping
            max_abs = np.max(np.abs(sig))
            if max_abs > 1e-4:
                sig = sig * (0.95 / max_abs)

            # Tandem Stage 1: Cellular GSM-FR 6.10
            gsm_out = self._compress_gsm_single(sig)
            # Ensure exact length alignment
            if len(gsm_out) > target_len:
                gsm_out = gsm_out[:target_len]
            elif len(gsm_out) < target_len:
                gsm_out = np.pad(gsm_out, (0, target_len - len(gsm_out)))

            # Tandem Stage 2: VoIP G.711 µ-law or Opus
            voip_out = self._compress_voip_single(gsm_out)
            if len(voip_out) > target_len:
                voip_out = voip_out[:target_len]
            elif len(voip_out) < target_len:
                voip_out = np.pad(voip_out, (0, target_len - len(voip_out)))

            processed_batch.append(voip_out)

        out_tensor = torch.from_numpy(np.stack(processed_batch)).to(device=x.device, dtype=x.dtype)
        out_tensor = out_tensor.reshape(*orig_shape[:-1], target_len)

        # Step 5: VoIP Packet loss simulation
        out_tensor = self.packet_loss(out_tensor)

        # Step 6: Peak normalization
        max_val = torch.max(torch.abs(out_tensor), dim=-1, keepdim=True)[0]
        max_val = torch.clamp(max_val, min=1e-6)
        scale = torch.where(max_val > 0.01, 0.95 / max_val, torch.ones_like(max_val))
        out_tensor = out_tensor * scale

        return out_tensor