"""Advanced Physics and Mathematical Feature Extractors for Telephony Audio (0-4 kHz).

Implements:
1. Teager-Kaiser Energy Operator (TKEO): Non-linear mechanical energy tracking.
2. Glottal Inverse Filtering (GIF): Vocal tract removal, glottal residual isolation,
   and cycle-to-cycle Jitter and Shimmer micro-perturbation analysis.
3. Modified Group Delay (MGD): High-resolution phase continuity representations.
4. Bicoherence (Higher-Order Spectra): Vectorized Quadratic Phase Coupling (QPC).
"""

from typing import Dict, Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F

from src.features import compute_autocorr, levinson_durbin


class TeagerKaiserEnergyExtractor(nn.Module):
    """Teager-Kaiser Energy Operator (TKEO) module.
    
    Discrete non-linear energy tracking:
        ψ[x[n]] = x[n]^2 - x[n+1] * x[n-1]
    
    Extracts frame-level energy statistics:
    1. Mean Teager Energy
    2. Teager Energy Variance
    3. Crest Factor (Peak-to-Average Ratio)
    """

    def __init__(self, frame_len: int = 200, hop_len: int = 80):
        super().__init__()
        self.frame_len = frame_len
        self.hop_len = hop_len

    def forward(self, waveform: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Args:
            waveform: Tensor of shape (B, N) at 8 kHz.
        Returns:
            raw_tkeo: (B, N) instantaneous TKEO sequence.
            frame_stats: (B, 3, num_frames) [mean, var, crest_factor].
        """
        # 1. Compute sample-wise TKEO: ψ[n] = x[n]^2 - x[n+1]*x[n-1]
        x = waveform
        tkeo_inner = x[..., 1:-1] ** 2 - x[..., 2:] * x[..., :-2]
        # Pad edges to maintain exact length N
        raw_tkeo = F.pad(tkeo_inner, (1, 1), mode="replicate")

        # 2. Frame the TKEO sequence: (B, num_frames, frame_len)
        frames = raw_tkeo.unfold(dimension=-1, size=self.frame_len, step=self.hop_len)

        # 3. Compute frame statistics
        mean_energy = torch.mean(frames, dim=-1)
        var_energy = torch.var(frames, dim=-1, unbiased=False)
        max_energy = torch.max(frames, dim=-1)[0]
        crest_factor = max_energy / torch.clamp(torch.abs(mean_energy), min=1e-6)

        # Stack: (B, 3, num_frames)
        frame_stats = torch.stack([mean_energy, var_energy, crest_factor], dim=1)
        return raw_tkeo, frame_stats


class GlottalInverseFilterExtractor(nn.Module):
    """Glottal Inverse Filtering (GIF) with Jitter and Shimmer extraction.
    
    Removes vocal tract formants via LPC inverse filtering to isolate the glottal source residual e[n].
    Measures pitch perturbation (Jitter) and amplitude perturbation (Shimmer) across glottal cycles.
    """

    def __init__(
        self,
        sample_rate: int = 8000,
        frame_len: int = 200,
        hop_len: int = 80,
        lpc_order: int = 14,
        min_pitch_hz: float = 60.0,
        max_pitch_hz: float = 400.0,
    ):
        super().__init__()
        self.sample_rate = sample_rate
        self.frame_len = frame_len
        self.hop_len = hop_len
        self.lpc_order = lpc_order

        # Pitch search range in samples at 8 kHz (e.g. 60 Hz -> 133 samples, 400 Hz -> 20 samples)
        self.min_lag = max(2, int(sample_rate / max_pitch_hz))
        self.max_lag = min(frame_len - 1, int(sample_rate / min_pitch_hz))

    def forward(self, waveform: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Args:
            waveform: Tensor of shape (B, N) at 8 kHz.
        Returns:
            residual: (B, N) glottal excitation waveform e[n].
            glottal_stats: (B, 3, num_frames) [residual_rms, local_jitter, local_shimmer].
        """
        B, N = waveform.shape

        # 1. Pre-emphasis to cancel lip radiation (1 - 0.97 z^-1)
        pre = torch.cat([waveform[..., :1], waveform[..., 1:] - 0.97 * waveform[..., :-1]], dim=-1)

        # 2. Frame signal and estimate vocal tract all-pole filter via Levinson-Durbin
        frames = pre.unfold(dimension=-1, size=self.frame_len, step=self.hop_len)
        r = compute_autocorr(frames, self.lpc_order)
        lpc, _ = levinson_durbin(r, self.lpc_order)  # (B, num_frames, lpc_order)

        # 3. Time-domain inverse filtering e[n] = x[n] - sum(a_k * x[n-k])
        # Compute global/mean LPC filter per sample or frame-wise
        # Mean LPC filter across frames for stable global inverse filtering
        mean_lpc = torch.mean(lpc, dim=1)  # (B, lpc_order)

        # Form FIR inverse filter: [1, -a_1, -a_2, ..., -a_P]
        # In conv1d, kernel is (out_channels, in_channels, kernel_size)
        residuals = []
        for b in range(B):
            inv_filter = torch.cat([torch.tensor([1.0], device=waveform.device), -mean_lpc[b]])
            inv_kernel = inv_filter.flip(dims=[0]).view(1, 1, -1)
            sig_b = pre[b:b+1].unsqueeze(1)
            pad_left = self.lpc_order
            sig_padded = F.pad(sig_b, (pad_left, 0), mode="replicate")
            res_b = F.conv1d(sig_padded, inv_kernel).squeeze(1)  # (1, N)
            residuals.append(res_b)

        residual = torch.cat(residuals, dim=0)  # (B, N)

        # 4. Measure Jitter and Shimmer on framed glottal residual
        res_frames = residual.unfold(dimension=-1, size=self.frame_len, step=self.hop_len)

        # Autocorrelation of residual frames for pitch cycle tracking
        n_fft = 2 ** ((2 * self.frame_len - 1).bit_length())
        res_fft = torch.fft.rfft(res_frames, n=n_fft, dim=-1)
        res_r = torch.fft.irfft(res_fft.abs() ** 2, n=n_fft, dim=-1)[..., :self.frame_len]

        # Search pitch lag in [min_lag, max_lag]
        pitch_lags = res_r[..., self.min_lag:self.max_lag + 1]
        best_lag_rel = torch.argmax(pitch_lags, dim=-1)
        t0 = best_lag_rel.float() + self.min_lag  # Pitch period in samples (B, num_frames)

        # Amplitude peak per frame
        a0 = torch.max(torch.abs(res_frames), dim=-1)[0]  # (B, num_frames)

        # Frame residual RMS
        res_rms = torch.sqrt(torch.mean(res_frames ** 2, dim=-1) + 1e-8)

        # Local Jitter: cycle-to-cycle period perturbation |T0[t] - T0[t-1]| / mean(T0)
        diff_t0 = torch.abs(t0[:, 1:] - t0[:, :-1])
        diff_t0 = torch.cat([diff_t0[:, :1], diff_t0], dim=-1)
        mean_t0 = torch.clamp(torch.mean(t0, dim=-1, keepdim=True), min=1.0)
        jitter = diff_t0 / mean_t0

        # Local Shimmer: cycle-to-cycle amplitude perturbation |A0[t] - A0[t-1]| / mean(A0)
        diff_a0 = torch.abs(a0[:, 1:] - a0[:, :-1])
        diff_a0 = torch.cat([diff_a0[:, :1], diff_a0], dim=-1)
        mean_a0 = torch.clamp(torch.mean(a0, dim=-1, keepdim=True), min=1e-5)
        shimmer = diff_a0 / mean_a0

        # Stack: (B, 3, num_frames)
        glottal_stats = torch.stack([res_rms, jitter, shimmer], dim=1)
        return residual, glottal_stats


class ModifiedGroupDelayExtractor(nn.Module):
    """Modified Group Delay (MGD) feature extractor.
    
    Computes negative derivative of unwrapped phase spectrum using Fourier transforms:
        τ(ω) = (X_R * Y_R + X_I * Y_I) / S(ω)^(2γ)
        MGD(ω) = sign(τ) * |τ|^α
    
    Exposes unnatural vocoder phase continuity and phase unwrapping irregularities.
    """

    def __init__(
        self,
        n_fft: int = 256,
        frame_len: int = 200,
        hop_len: int = 80,
        gamma: float = 0.9,
        alpha: float = 0.4,
        num_ceps: int = 64
    ):
        super().__init__()
        self.n_fft = n_fft
        self.frame_len = frame_len
        self.hop_len = hop_len
        self.gamma = gamma
        self.alpha = alpha
        self.num_ceps = num_ceps

        # Ramp vector n: [0, 1, 2, ..., frame_len - 1]
        self.register_buffer("ramp", torch.arange(frame_len, dtype=torch.float32).view(1, 1, -1).contiguous().clone(), persistent=False)

    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
        """Args:
            waveform: Tensor of shape (B, N) at 8 kHz.
        Returns:
            mgd: Tensor of shape (B, 1, num_ceps, num_frames)
        """
        # 1. Frame signal
        x_frames = waveform.unfold(dimension=-1, size=self.frame_len, step=self.hop_len)  # (B, T, L)
        y_frames = x_frames * self.ramp  # n * x[n]

        # 2. FFT of x[n] and n * x[n]
        X = torch.fft.rfft(x_frames, n=self.n_fft, dim=-1)  # (B, T, F)
        Y = torch.fft.rfft(y_frames, n=self.n_fft, dim=-1)

        # 3. Numerator of group delay: X_R * Y_R + X_I * Y_I
        num = X.real * Y.real + X.imag * Y.imag

        # 4. Smoothed spectral magnitude envelope S(ω) via moving average across frequency
        mag = X.abs()
        B_dim, T_dim, F_dim = mag.shape
        mag_flat = mag.reshape(-1, 1, F_dim)
        smooth_mag = F.avg_pool1d(mag_flat, kernel_size=5, stride=1, padding=2).reshape(B_dim, T_dim, F_dim)
        denom = torch.clamp(smooth_mag, min=1e-5) ** (2.0 * self.gamma)

        # 5. Modified Group Delay
        tau = num / denom
        mgd = torch.sign(tau) * (torch.clamp(torch.abs(tau), min=1e-8) ** self.alpha)

        # Select first num_ceps frequency bins (or interpolate)
        if mgd.shape[-1] >= self.num_ceps:
            mgd = mgd[..., :self.num_ceps]
        else:
            mgd = F.pad(mgd, (0, self.num_ceps - mgd.shape[-1]))

        # Permute to (B, 1, num_ceps, num_frames) for 2D CNN
        mgd_2d = mgd.transpose(1, 2).unsqueeze(1)
        return mgd_2d


class VectorizedBicoherenceExtractor(nn.Module):
    """Vectorized Bicoherence (Higher-Order Spectra / Bispectrum) module.
    
    Computes Quadratic Phase Coupling (QPC) in the 0-4 kHz band:
        B(f1, f2) = E[X(f1) * X(f2) * X*(f1 + f2)]
        b^2(f1, f2) = |B(f1, f2)|^2 / (E[|X(f1)X(f2)|^2] * E[|X(f1+f2)|^2])
    
    Reveals lack of non-linear physical harmonic locking in synthetic neural vocoders.
    """

    def __init__(self, n_fft: int = 128, frame_len: int = 200, hop_len: int = 80, num_bins: int = 32):
        super().__init__()
        self.n_fft = n_fft
        self.frame_len = frame_len
        self.hop_len = hop_len
        self.num_bins = num_bins

        # Precompute non-redundant frequency indices grid (i, j) where i + j < num_bins and j >= i
        i_idx, j_idx = torch.meshgrid(torch.arange(num_bins), torch.arange(num_bins), indexing="ij")
        mask = (i_idx + j_idx < num_bins) & (j_idx >= i_idx)
        sum_idx = torch.clamp(i_idx + j_idx, max=num_bins - 1)

        self.register_buffer("i_idx", i_idx.contiguous().clone(), persistent=False)
        self.register_buffer("j_idx", j_idx.contiguous().clone(), persistent=False)
        self.register_buffer("sum_idx", sum_idx.contiguous().clone(), persistent=False)
        self.register_buffer("mask", mask.contiguous().clone(), persistent=False)

    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
        """Args:
            waveform: Tensor of shape (B, N) at 8 kHz.
        Returns:
            bicoherence: Tensor of shape (B, 1, num_bins, num_bins)
        """
        # 1. Frame signal
        frames = waveform.unfold(dimension=-1, size=self.frame_len, step=self.hop_len)  # (B, T, L)

        # 2. Window and FFT
        window = torch.hamming_window(self.frame_len, device=waveform.device)
        w_frames = frames * window
        X = torch.fft.rfft(w_frames, n=self.n_fft, dim=-1)  # (B, T, F_fft)

        # Select first num_bins frequency channels (0 - 4 kHz band)
        X_sub = X[..., :self.num_bins]  # (B, T, num_bins) complex64

        # 3. Vectorized tensor index expansion
        X_i = X_sub[:, :, self.i_idx]      # (B, T, num_bins, num_bins)
        X_j = X_sub[:, :, self.j_idx]      # (B, T, num_bins, num_bins)
        X_sum = X_sub[:, :, self.sum_idx]  # (B, T, num_bins, num_bins)

        # Bispectrum expectation across frames dimension (T)
        term = X_i * X_j * torch.conj(X_sum)
        bispectrum = torch.mean(term, dim=1)  # (B, num_bins, num_bins)

        # Normalization denominator
        denom1 = torch.mean((X_i * X_j).abs() ** 2, dim=1)
        denom2 = torch.mean(X_sum.abs() ** 2, dim=1)
        denom = torch.clamp(denom1 * denom2, min=1e-8)

        # Normalized bicoherence
        bico = (bispectrum.abs() ** 2) / denom
        bico = bico * self.mask.unsqueeze(0).float()

        # Add channel dimension: (B, 1, num_bins, num_bins)
        return bico.unsqueeze(1)


class StackedPhysicsFeatureExtractor(nn.Module):
    """Orchestrates all physics and mathematical feature extractors.
    
    Combines:
    - 1D Temporal Features (TKEO stats + Glottal residual stats): (B, 6, num_frames)
    - 2D Phase Spectrum (Modified Group Delay MGD): (B, 1, 64, num_frames)
    - 2D Harmonic Coupling (Vectorized Bicoherence): (B, 1, 32, 32)
    """

    def __init__(
        self,
        sample_rate: int = 8000,
        frame_len: int = 200,
        hop_len: int = 80,
        lpc_order: int = 14,
        mgd_ceps: int = 64,
        bico_bins: int = 32
    ):
        super().__init__()
        self.tkeo = TeagerKaiserEnergyExtractor(frame_len=frame_len, hop_len=hop_len)
        self.gif = GlottalInverseFilterExtractor(
            sample_rate=sample_rate,
            frame_len=frame_len,
            hop_len=hop_len,
            lpc_order=lpc_order
        )
        self.mgd = ModifiedGroupDelayExtractor(
            frame_len=frame_len,
            hop_len=hop_len,
            num_ceps=mgd_ceps
        )
        self.bico = VectorizedBicoherenceExtractor(
            frame_len=frame_len,
            hop_len=hop_len,
            num_bins=bico_bins
        )

    def forward(self, waveform: torch.Tensor) -> Dict[str, torch.Tensor]:
        """Args:
            waveform: Tensor of shape (B, N) at 8 kHz.
        Returns:
            Dict containing:
            - 'feats_1d': (B, 6, num_frames) [TKEO stats (3) + Glottal stats (3)]
            - 'mgd':      (B, 1, 64, num_frames) [Modified Group Delay 2D]
            - 'bico':     (B, 1, 32, 32) [Bicoherence Matrix 2D]
            - 'raw_tkeo': (B, N)
            - 'residual': (B, N)
        """
        raw_tkeo, tkeo_stats = self.tkeo(waveform)
        residual, gif_stats = self.gif(waveform)
        mgd_2d = self.mgd(waveform)
        bico_2d = self.bico(waveform)

        # Align 1D temporal frame counts if minor boundary discrepancy
        min_frames = min(tkeo_stats.shape[-1], gif_stats.shape[-1], mgd_2d.shape[-1])
        tkeo_stats = tkeo_stats[..., :min_frames]
        gif_stats = gif_stats[..., :min_frames]
        mgd_2d = mgd_2d[..., :min_frames]

        # Stack 1D temporal physics features: (B, 6, min_frames)
        feats_1d = torch.cat([tkeo_stats, gif_stats], dim=1)

        return {
            "feats_1d": feats_1d,
            "mgd": mgd_2d,
            "bico": bico_2d,
            "raw_tkeo": raw_tkeo,
            "residual": residual
        }