"""Hybrid Feature Extraction Module for ASVspoof 5 Telephony Audio.

Implements:
1. Branch A (Physics-Based Vocal Tract & Glottal Features, 0-4 kHz):
   - Linear Prediction Analysis (LPC) via Levinson-Durbin recursion.
   - Linear Prediction Cepstral Coefficients (LPCC) capturing the all-pole vocal tract transfer function.
   - Glottal inverse filtering residual energy and excitation constraints.

2. Branch B (Spectrogram / Filterbank Features, 0-4 kHz):
   - Log-Mel Spectrogram and Linear Frequency Cepstral Coefficients (LFCC).
"""

from typing import Dict, Optional, Tuple
import math
import torch
import torch.nn as nn
import torchaudio.transforms as T


def compute_autocorr(frames: torch.Tensor, order: int) -> torch.Tensor:
    """Compute autocorrelation for framed signals using FFT.
    
    Args:
        frames: Tensor of shape (..., num_frames, frame_len)
        order: Maximum lag to compute (P)
    Returns:
        Autocorrelation tensor of shape (..., num_frames, order + 1)
    """
    frame_len = frames.shape[-1]
    window = torch.hamming_window(frame_len, device=frames.device, dtype=frames.dtype)
    w_frames = frames * window

    n_fft = 2 ** ((2 * frame_len - 1).bit_length())
    freq = torch.fft.rfft(w_frames, n=n_fft, dim=-1)
    power = freq.abs() ** 2
    r = torch.fft.irfft(power, n=n_fft, dim=-1)[..., :order + 1]
    return r


def levinson_durbin(r: torch.Tensor, order: int, eps: float = 1e-8) -> Tuple[torch.Tensor, torch.Tensor]:
    """Vectorized Levinson-Durbin algorithm.
    
    Solves R * a = -r for all-pole vocal tract filter coefficients.
    
    Args:
        r: Autocorrelation tensor of shape (..., order + 1)
        order: LPC order P
    Returns:
        a: LPC filter coefficients of shape (..., order)
        e: Prediction error / residual power of shape (...)
    """
    batch_shape = r.shape[:-1]
    r_flat = r.reshape(-1, order + 1)
    B = r_flat.shape[0]

    a = torch.zeros(B, order + 1, device=r.device, dtype=r.dtype)
    a[:, 0] = 1.0

    e = torch.clamp(r_flat[:, 0].clone(), min=eps)

    for i in range(1, order + 1):
        acc = torch.zeros(B, device=r.device, dtype=r.dtype)
        for j in range(1, i):
            acc += a[:, j] * r_flat[:, i - j]
        k = -(r_flat[:, i] + acc) / e

        # Reflection coefficient stability clamp
        k = torch.clamp(k, -0.999, 0.999)

        new_a = a.clone()
        for j in range(1, i):
            new_a[:, j] = a[:, j] + k * a[:, i - j]
        new_a[:, i] = k
        a = new_a

        e = e * (1.0 - k ** 2)
        e = torch.clamp(e, min=eps)

    lpc_coeffs = a[:, 1:].reshape(*batch_shape, order)
    residual_power = e.reshape(*batch_shape)
    return lpc_coeffs, residual_power


def lpc_to_lpcc(lpc: torch.Tensor, num_cepstrum: int) -> torch.Tensor:
    """Recursively convert LPC coefficients to Linear Prediction Cepstral Coefficients (LPCC).
    
    Captures the cepstral envelope of the vocal tract transfer function.
    """
    *batch_shape, P = lpc.shape
    lpc_flat = lpc.reshape(-1, P)
    B = lpc_flat.shape[0]

    cep = torch.zeros(B, num_cepstrum, device=lpc.device, dtype=lpc.dtype)
    cep[:, 0] = -lpc_flat[:, 0]

    for n in range(1, num_cepstrum):
        idx = n + 1
        sum_val = torch.zeros(B, device=lpc.device, dtype=lpc.dtype)
        limit = min(n, P)
        for k in range(limit):
            k_1idx = k + 1
            sum_val += (1.0 - k_1idx / idx) * lpc_flat[:, k] * cep[:, n - k_1idx]

        if n < P:
            cep[:, n] = -lpc_flat[:, n] - sum_val
        else:
            cep[:, n] = -sum_val

    return cep.reshape(*batch_shape, num_cepstrum)


class PhysicsVocalTractExtractor(nn.Module):
    """Branch A: Physics-based Vocal Tract & Glottal Feature Extractor (0-4 kHz).
    
    Extracts:
    1. Linear Prediction Cepstral Coefficients (LPCC): vocal tract resonance shape.
    2. Glottal residual energy ratio & predictability: glottal excitation constraints.
    """

    def __init__(
        self,
        sample_rate: int = 8000,
        frame_len: int = 200,   # 25 ms @ 8 kHz
        hop_len: int = 80,      # 10 ms @ 8 kHz
        lpc_order: int = 14,    # Appropriate for 8kHz narrowband speech
        num_lpcc: int = 16
    ):
        super().__init__()
        self.sample_rate = sample_rate
        self.frame_len = frame_len
        self.hop_len = hop_len
        self.lpc_order = lpc_order
        self.num_lpcc = num_lpcc
        self.out_channels = num_lpcc + 2  # LPCC + log residual power + residual ratio

    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
        """Extract physics-based features.
        
        Args:
            waveform: (B, num_samples) at 8 kHz
        Returns:
            features: (B, out_channels, num_frames)
        """
        # 1. Frame signal
        # (B, num_frames, frame_len)
        frames = waveform.unfold(dimension=-1, size=self.frame_len, step=self.hop_len)

        # 2. Compute Autocorrelation
        r = compute_autocorr(frames, self.lpc_order)

        # 3. Levinson-Durbin LPC & Residual Power
        lpc, res_power = levinson_durbin(r, self.lpc_order)

        # 4. Convert LPC to LPCC
        lpcc = lpc_to_lpcc(lpc, self.num_lpcc)  # (B, num_frames, num_lpcc)

        # 5. Glottal residual constraints
        # Frame signal energy
        frame_energy = torch.clamp(r[..., 0], min=1e-8)
        log_res_power = torch.log(torch.clamp(res_power, min=1e-8)).unsqueeze(-1)
        res_ratio = (res_power / frame_energy).unsqueeze(-1)  # Predictability ratio

        # Combine: (B, num_frames, num_lpcc + 2)
        phys_features = torch.cat([lpcc, log_res_power, res_ratio], dim=-1)

        # Permute to (B, channels, time) for 1D convolutions
        return phys_features.transpose(1, 2)


class SpectrogramExtractor(nn.Module):
    """Branch B: Spectrogram & Filterbank Extractor (0-4 kHz).
    
    Supports Log-Mel Spectrogram and Linear Frequency Cepstral Coefficients (LFCC).
    """

    def __init__(
        self,
        sample_rate: int = 8000,
        n_fft: int = 512,
        frame_len: int = 200,   # 25 ms @ 8 kHz
        hop_len: int = 80,      # 10 ms @ 8 kHz
        n_mels: int = 64,
        f_min: float = 0.0,
        f_max: float = 4000.0,
        feature_type: str = "mel"  # "mel" or "lfcc"
    ):
        super().__init__()
        self.sample_rate = sample_rate
        self.feature_type = feature_type
        self.n_mels = n_mels

        self.mel_transform = T.MelSpectrogram(
            sample_rate=sample_rate,
            n_fft=n_fft,
            win_length=frame_len,
            hop_length=hop_len,
            n_mels=n_mels,
            f_min=f_min,
            f_max=f_max,
            power=2.0
        )

    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
        """Extract spectral features.
        
        Args:
            waveform: (B, num_samples) at 8 kHz
        Returns:
            features: (B, 1, freq_bins, time_frames)
        """
        # Mel spectrogram: (B, n_mels, T)
        mel = self.mel_transform(waveform)
        log_mel = torch.log(torch.clamp(mel, min=1e-6))
        # Add channel dimension for 2D CNN: (B, 1, n_mels, T)
        return log_mel.unsqueeze(1)


class HybridFeatureExtractor(nn.Module):
    """Dual-branch hybrid feature extractor orchestrating physics and spectral branches."""

    def __init__(
        self,
        sample_rate: int = 8000,
        frame_len: int = 200,
        hop_len: int = 80,
        lpc_order: int = 14,
        num_lpcc: int = 16,
        n_mels: int = 64
    ):
        super().__init__()
        self.branch_a = PhysicsVocalTractExtractor(
            sample_rate=sample_rate,
            frame_len=frame_len,
            hop_len=hop_len,
            lpc_order=lpc_order,
            num_lpcc=num_lpcc
        )
        self.branch_b = SpectrogramExtractor(
            sample_rate=sample_rate,
            frame_len=frame_len,
            hop_len=hop_len,
            n_mels=n_mels
        )

    def forward(self, waveform: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Extract both Branch A and Branch B features.
        
        Args:
            waveform: (B, num_samples) at 8 kHz
        Returns:
            phys_features: (B, 18, T)
            spec_features: (B, 1, 64, T)
        """
        phys = self.branch_a(waveform)
        spec = self.branch_b(waveform)
        return phys, spec