rustytorch/crates/specialized/rtx-neuro-connectivity/src/pac.rs

//! Phase-Amplitude Coupling (PAC) analysis.
//!
//! PAC measures the relationship between the phase of low-frequency oscillations
//! and the amplitude of high-frequency oscillations.
//!
//! ## Mathematical Background
//!
//! Modulation Index (MI): MI = KL(P || U) / log(N)
//! where P is the distribution of amplitudes over phase bins and U is uniform.

use crate::{ConnectivityError, ConnectivityResult, utils};
use num_complex::Complex64;
use std::f64::consts::PI;

/// PAC estimation method
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PacMethod {
    /// Modulation Index (Tort et al.)
    ModulationIndex,
    /// Mean Vector Length
    MeanVectorLength,
    /// Phase-Locking Value between phase and amplitude envelope
    PlvPac,
}

impl Default for PacMethod {
    fn default() -> Self {
        Self::ModulationIndex
    }
}

/// PAC result for a comodulogram
#[derive(Debug, Clone)]
pub struct PacResult {
    /// PAC values [n_phase_freqs x n_amp_freqs]
    pub data: Vec<Vec<f64>>,
    /// Phase frequencies
    pub phase_freqs: Vec<f64>,
    /// Amplitude frequencies
    pub amp_freqs: Vec<f64>,
    /// Method used
    pub method: PacMethod,
}

/// Compute Phase-Amplitude Coupling
///
/// # Arguments
/// * `signal` - Time series [n_samples]
/// * `sfreq` - Sampling frequency
/// * `phase_freq` - (low, high) frequency band for phase
/// * `amp_freq` - (low, high) frequency band for amplitude
/// * `method` - PAC estimation method
///
/// # Returns
/// PAC value
pub fn phase_amplitude_coupling(
    signal: &[f64],
    sfreq: f64,
    phase_freq: (f64, f64),
    amp_freq: (f64, f64),
    method: PacMethod,
) -> ConnectivityResult<f64> {
    if signal.len() < 100 {
        return Err(ConnectivityError::InsufficientData(
            "Signal too short for PAC analysis".to_string(),
        ));
    }

    if phase_freq.1 >= amp_freq.0 {
        return Err(ConnectivityError::InvalidParameters(
            "Phase frequency band must be lower than amplitude frequency band".to_string(),
        ));
    }

    // Bandpass filter for phase and amplitude
    let phase_signal = utils::bandpass_fft(signal, sfreq, phase_freq.0, phase_freq.1);
    let amp_signal = utils::bandpass_fft(signal, sfreq, amp_freq.0, amp_freq.1);

    // Get phase from low-frequency signal
    let phase_analytic = utils::hilbert(&phase_signal);
    let phase = utils::instantaneous_phase(&phase_analytic);

    // Get amplitude envelope from high-frequency signal
    let amp_analytic = utils::hilbert(&amp_signal);
    let amplitude: Vec<f64> = amp_analytic.iter().map(|c| c.norm()).collect();

    match method {
        PacMethod::ModulationIndex => modulation_index(&phase, &amplitude),
        PacMethod::MeanVectorLength => mean_vector_length(&phase, &amplitude),
        PacMethod::PlvPac => plv_pac(&phase, &amplitude),
    }
}

/// Compute comodulogram (PAC across multiple frequency pairs)
///
/// # Arguments
/// * `signal` - Time series
/// * `sfreq` - Sampling frequency
/// * `phase_freqs` - Center frequencies for phase (low frequency)
/// * `amp_freqs` - Center frequencies for amplitude (high frequency)
/// * `bandwidth` - Bandwidth for bandpass filters
/// * `method` - PAC method
pub fn comodulogram(
    signal: &[f64],
    sfreq: f64,
    phase_freqs: &[f64],
    amp_freqs: &[f64],
    bandwidth: f64,
    method: PacMethod,
) -> ConnectivityResult<PacResult> {
    let mut data = Vec::with_capacity(phase_freqs.len());

    for &f_phase in phase_freqs {
        let mut row = Vec::with_capacity(amp_freqs.len());

        for &f_amp in amp_freqs {
            // Ensure non-overlapping frequency bands
            if f_phase + bandwidth / 2.0 >= f_amp - bandwidth / 2.0 {
                row.push(0.0);
                continue;
            }

            let phase_band = (f_phase - bandwidth / 2.0, f_phase + bandwidth / 2.0);
            let amp_band = (f_amp - bandwidth / 2.0, f_amp + bandwidth / 2.0);

            let pac = phase_amplitude_coupling(signal, sfreq, phase_band, amp_band, method)
                .unwrap_or(0.0);
            row.push(pac);
        }

        data.push(row);
    }

    Ok(PacResult {
        data,
        phase_freqs: phase_freqs.to_vec(),
        amp_freqs: amp_freqs.to_vec(),
        method,
    })
}

/// Modulation Index (Tort et al., 2010)
///
/// Measures deviation from uniform distribution of amplitudes across phase bins
fn modulation_index(phase: &[f64], amplitude: &[f64]) -> ConnectivityResult<f64> {
    let n_bins = 18; // 20-degree bins
    let bin_width = 2.0 * PI / n_bins as f64;

    // Bin amplitudes by phase
    let mut bin_sums = vec![0.0; n_bins];
    let mut bin_counts = vec![0usize; n_bins];

    for (&p, &a) in phase.iter().zip(amplitude) {
        // Normalize phase to [0, 2π)
        let p_norm = ((p % (2.0 * PI)) + 2.0 * PI) % (2.0 * PI);
        let bin = ((p_norm / bin_width) as usize).min(n_bins - 1);
        bin_sums[bin] += a;
        bin_counts[bin] += 1;
    }

    // Compute mean amplitude per bin
    let bin_means: Vec<f64> = bin_sums
        .iter()
        .zip(&bin_counts)
        .map(|(&s, &c)| if c > 0 { s / c as f64 } else { 0.0 })
        .collect();

    // Normalize to probability distribution
    let total: f64 = bin_means.iter().sum();
    if total < 1e-15 {
        return Ok(0.0);
    }

    let p_dist: Vec<f64> = bin_means.iter().map(|&m| m / total).collect();

    // KL divergence from uniform distribution
    let uniform = 1.0 / n_bins as f64;
    let mut kl = 0.0;
    for &p in &p_dist {
        if p > 1e-15 {
            kl += p * (p / uniform).ln();
        }
    }

    // Normalize by log(N) to get MI in [0, 1]
    let mi = kl / (n_bins as f64).ln();
    Ok(mi.max(0.0))
}

/// Mean Vector Length (MVL)
///
/// Measures the length of the mean vector of amplitude-weighted phase vectors
fn mean_vector_length(phase: &[f64], amplitude: &[f64]) -> ConnectivityResult<f64> {
    let n = phase.len();
    if n == 0 {
        return Ok(0.0);
    }

    // Compute amplitude-weighted mean vector
    let mut sum_cos = 0.0;
    let mut sum_sin = 0.0;
    let mut sum_amp = 0.0;

    for (&p, &a) in phase.iter().zip(amplitude) {
        sum_cos += a * p.cos();
        sum_sin += a * p.sin();
        sum_amp += a;
    }

    if sum_amp < 1e-15 {
        return Ok(0.0);
    }

    let mean_cos = sum_cos / sum_amp;
    let mean_sin = sum_sin / sum_amp;

    let mvl = (mean_cos * mean_cos + mean_sin * mean_sin).sqrt();
    Ok(mvl)
}

/// PLV-based PAC
///
/// Computes PLV between low-frequency phase and high-frequency amplitude envelope
fn plv_pac(phase: &[f64], amplitude: &[f64]) -> ConnectivityResult<f64> {
    let n = phase.len();
    if n == 0 {
        return Ok(0.0);
    }

    // Get phase of amplitude envelope (via Hilbert)
    let amp_analytic = utils::hilbert(amplitude);
    let amp_phase: Vec<f64> = amp_analytic.iter().map(|c| c.arg()).collect();

    // Compute PLV between phases
    let mut sum = Complex64::new(0.0, 0.0);
    for (&p1, &p2) in phase.iter().zip(&amp_phase) {
        let diff = p1 - p2;
        sum += Complex64::new(diff.cos(), diff.sin());
    }

    let plv = sum.norm() / n as f64;
    Ok(plv)
}

/// Compute PAC for epoched data
pub fn pac_epochs(
    epochs: &[Vec<f64>],
    sfreq: f64,
    phase_freq: (f64, f64),
    amp_freq: (f64, f64),
    method: PacMethod,
) -> ConnectivityResult<f64> {
    if epochs.is_empty() {
        return Err(ConnectivityError::InsufficientData(
            "No epochs provided".to_string(),
        ));
    }

    // Average PAC across epochs
    let mut pac_sum = 0.0;
    let mut count = 0;

    for epoch in epochs {
        match phase_amplitude_coupling(epoch, sfreq, phase_freq, amp_freq, method) {
            Ok(pac) => {
                pac_sum += pac;
                count += 1;
            }
            Err(_) => continue,
        }
    }

    if count > 0 {
        Ok(pac_sum / count as f64)
    } else {
        Err(ConnectivityError::ComputationError(
            "Failed to compute PAC for any epoch".to_string(),
        ))
    }
}

/// Surrogate-based statistical testing for PAC
///
/// Shuffles amplitude time series to create null distribution
pub fn pac_permutation_test(
    signal: &[f64],
    sfreq: f64,
    phase_freq: (f64, f64),
    amp_freq: (f64, f64),
    method: PacMethod,
    n_permutations: usize,
) -> ConnectivityResult<(f64, f64, f64)> {
    // Compute observed PAC
    let pac_obs = phase_amplitude_coupling(signal, sfreq, phase_freq, amp_freq, method)?;

    // Get phase and amplitude time series
    let phase_signal = utils::bandpass_fft(signal, sfreq, phase_freq.0, phase_freq.1);
    let amp_signal = utils::bandpass_fft(signal, sfreq, amp_freq.0, amp_freq.1);

    let phase_analytic = utils::hilbert(&phase_signal);
    let phase = utils::instantaneous_phase(&phase_analytic);

    let amp_analytic = utils::hilbert(&amp_signal);
    let amplitude: Vec<f64> = amp_analytic.iter().map(|c| c.norm()).collect();

    // Generate surrogate distribution
    let n = amplitude.len();
    let mut surrogate_pacs = Vec::with_capacity(n_permutations);

    for i in 0..n_permutations {
        // Circular shift of amplitude
        let shift = (n * (i + 1) / (n_permutations + 1)).max(n / 10);
        let amp_shifted: Vec<f64> = amplitude
            .iter()
            .cycle()
            .skip(shift)
            .take(n)
            .copied()
            .collect();

        let pac = match method {
            PacMethod::ModulationIndex => modulation_index(&phase, &amp_shifted),
            PacMethod::MeanVectorLength => mean_vector_length(&phase, &amp_shifted),
            PacMethod::PlvPac => plv_pac(&phase, &amp_shifted),
        };

        if let Ok(p) = pac {
            surrogate_pacs.push(p);
        }
    }

    if surrogate_pacs.is_empty() {
        return Err(ConnectivityError::ComputationError(
            "Failed to generate surrogates".to_string(),
        ));
    }

    // Compute p-value
    let n_greater = surrogate_pacs.iter().filter(|&&p| p >= pac_obs).count();
    let p_value = n_greater as f64 / surrogate_pacs.len() as f64;

    // Compute z-score
    let mean_surr: f64 = surrogate_pacs.iter().sum::<f64>() / surrogate_pacs.len() as f64;
    let var_surr: f64 = surrogate_pacs
        .iter()
        .map(|&p| (p - mean_surr) * (p - mean_surr))
        .sum::<f64>()
        / surrogate_pacs.len() as f64;
    let std_surr = var_surr.sqrt();

    let z_score = if std_surr > 1e-15 {
        (pac_obs - mean_surr) / std_surr
    } else {
        0.0
    };

    Ok((pac_obs, p_value, z_score))
}

#[cfg(test)]
mod tests {
    use super::*;

    fn create_pac_signal(
        n: usize,
        sfreq: f64,
        phase_freq: f64,
        amp_freq: f64,
        coupling: f64,
    ) -> Vec<f64> {
        (0..n)
            .map(|i| {
                let t = i as f64 / sfreq;
                // Low frequency carrier
                let phase = 2.0 * PI * phase_freq * t;
                // High frequency with amplitude modulated by low frequency phase
                let amp_mod = 1.0 + coupling * phase.cos();
                let hf = amp_mod * (2.0 * PI * amp_freq * t).sin();
                phase.sin() + hf
            })
            .collect()
    }

    #[test]
    fn test_pac_with_coupling() {
        let signal = create_pac_signal(2000, 500.0, 6.0, 60.0, 0.8);
        let pac = phase_amplitude_coupling(
            &signal,
            500.0,
            (4.0, 8.0),
            (50.0, 70.0),
            PacMethod::ModulationIndex,
        );

        assert!(pac.is_ok());
        let pac_val = pac.unwrap();
        assert!(
            pac_val > 0.05,
            "PAC with coupling should be detectable, got {}",
            pac_val
        );
    }

    #[test]
    fn test_pac_without_coupling() {
        // Pure noise - no coupling
        let signal: Vec<f64> = (0..2000)
            .map(|i| ((i as f64 * 0.1).sin() * 100.0).sin())
            .collect();

        let pac = phase_amplitude_coupling(
            &signal,
            500.0,
            (4.0, 8.0),
            (50.0, 70.0),
            PacMethod::ModulationIndex,
        );

        assert!(pac.is_ok());
        let pac_val = pac.unwrap();
        // Should be low without coupling
        assert!(pac_val < 0.3);
    }

    #[test]
    fn test_mvl() {
        let signal = create_pac_signal(2000, 500.0, 6.0, 60.0, 0.8);
        let pac = phase_amplitude_coupling(
            &signal,
            500.0,
            (4.0, 8.0),
            (50.0, 70.0),
            PacMethod::MeanVectorLength,
        );

        assert!(pac.is_ok());
        let pac_val = pac.unwrap();
        assert!(pac_val >= 0.0 && pac_val <= 1.0);
    }

    #[test]
    fn test_comodulogram() {
        let signal = create_pac_signal(2000, 500.0, 8.0, 80.0, 0.5);
        let phase_freqs: Vec<f64> = (4..=12).step_by(2).map(|f| f as f64).collect();
        let amp_freqs: Vec<f64> = (50..=100).step_by(10).map(|f| f as f64).collect();

        let result = comodulogram(
            &signal,
            500.0,
            &phase_freqs,
            &amp_freqs,
            4.0,
            PacMethod::ModulationIndex,
        );

        assert!(result.is_ok());
        let comod = result.unwrap();
        assert_eq!(comod.data.len(), phase_freqs.len());
        assert_eq!(comod.data[0].len(), amp_freqs.len());
    }

    #[test]
    fn test_pac_epochs() {
        let epochs: Vec<Vec<f64>> = (0..5)
            .map(|_| create_pac_signal(500, 250.0, 6.0, 60.0, 0.6))
            .collect();

        let pac = pac_epochs(
            &epochs,
            250.0,
            (4.0, 8.0),
            (50.0, 70.0),
            PacMethod::ModulationIndex,
        );
        assert!(pac.is_ok());
    }

    #[test]
    fn test_invalid_freq_bands() {
        let signal = create_pac_signal(1000, 500.0, 6.0, 60.0, 0.5);

        // Phase freq higher than amp freq - should error
        let result = phase_amplitude_coupling(
            &signal,
            500.0,
            (50.0, 60.0),
            (4.0, 8.0),
            PacMethod::ModulationIndex,
        );
        assert!(result.is_err());
    }
}