//! Weighted Phase Lag Index (wPLI) and debiased wPLI.
//!
//! wPLI is robust to volume conduction and noise-related phase synchronization
//! by focusing only on consistent phase leads/lags.
//!
//! ## Mathematical Background
//!
//! wPLI = |⟨|Im(Sxy)| * sign(Im(Sxy))⟩| / ⟨|Im(Sxy)|⟩
//!
//! where Sxy is the cross-spectral density and Im() takes the imaginary part.

use crate::{ConnectivityError, ConnectivityResult, utils};
use num_complex::Complex64;

/// Compute wPLI between two signals across epochs
///
/// # Arguments
/// * `x_epochs` - First signal epochs [n_epochs][n_samples]
/// * `y_epochs` - Second signal epochs [n_epochs][n_samples]
/// * `sfreq` - Sampling frequency
/// * `fmin` - Minimum frequency
/// * `fmax` - Maximum frequency
///
/// # Returns
/// wPLI value (0-1)
pub fn wpli(
    x_epochs: &[Vec<f64>],
    y_epochs: &[Vec<f64>],
    sfreq: f64,
    fmin: f64,
    fmax: f64,
) -> ConnectivityResult<f64> {
    if x_epochs.len() != y_epochs.len() {
        return Err(ConnectivityError::DimensionMismatch(
            "Number of epochs must match".to_string(),
        ));
    }

    if x_epochs.is_empty() {
        return Err(ConnectivityError::InsufficientData(
            "No epochs provided".to_string(),
        ));
    }

    let _n_epochs = x_epochs.len();
    let n_samples = x_epochs[0].len();
    let n_fft = n_samples;
    let freq_indices = utils::freq_indices(sfreq, n_fft, fmin, fmax);

    if freq_indices.is_empty() {
        return Err(ConnectivityError::InvalidParameters(
            "No frequencies in range".to_string(),
        ));
    }

    let window = utils::hanning_window(n_samples);

    // Accumulate imaginary parts of cross-spectrum
    let mut im_sum = 0.0;
    let mut abs_im_sum = 0.0;

    for (x_epoch, y_epoch) in x_epochs.iter().zip(y_epochs) {
        let x_win = utils::apply_window(x_epoch, &window);
        let y_win = utils::apply_window(y_epoch, &window);

        let x_fft = utils::rfft(&x_win, n_fft);
        let y_fft = utils::rfft(&y_win, n_fft);

        for &k in &freq_indices {
            let csd = x_fft[k] * y_fft[k].conj();
            let im = csd.im;
            im_sum += im;
            abs_im_sum += im.abs();
        }
    }

    // wPLI = |sum(Im)| / sum(|Im|)
    if abs_im_sum > 1e-15 {
        Ok(im_sum.abs() / abs_im_sum)
    } else {
        Ok(0.0)
    }
}

/// Compute debiased wPLI
///
/// dwPLI corrects for sample size bias in wPLI
/// dwPLI = (sum(Im)² - sum(Im²)) / (sum(|Im|)² - sum(Im²))
pub fn dwpli(
    x_epochs: &[Vec<f64>],
    y_epochs: &[Vec<f64>],
    sfreq: f64,
    fmin: f64,
    fmax: f64,
) -> ConnectivityResult<f64> {
    if x_epochs.len() != y_epochs.len() {
        return Err(ConnectivityError::DimensionMismatch(
            "Number of epochs must match".to_string(),
        ));
    }

    if x_epochs.is_empty() {
        return Err(ConnectivityError::InsufficientData(
            "No epochs provided".to_string(),
        ));
    }

    let n_samples = x_epochs[0].len();
    let n_fft = n_samples;
    let freq_indices = utils::freq_indices(sfreq, n_fft, fmin, fmax);

    if freq_indices.is_empty() {
        return Err(ConnectivityError::InvalidParameters(
            "No frequencies in range".to_string(),
        ));
    }

    let window = utils::hanning_window(n_samples);

    // Collect all imaginary parts
    let mut im_values: Vec<f64> = Vec::new();

    for (x_epoch, y_epoch) in x_epochs.iter().zip(y_epochs) {
        let x_win = utils::apply_window(x_epoch, &window);
        let y_win = utils::apply_window(y_epoch, &window);

        let x_fft = utils::rfft(&x_win, n_fft);
        let y_fft = utils::rfft(&y_win, n_fft);

        for &k in &freq_indices {
            let csd = x_fft[k] * y_fft[k].conj();
            im_values.push(csd.im);
        }
    }

    // Compute debiased wPLI
    let sum_im: f64 = im_values.iter().sum();
    let sum_abs_im: f64 = im_values.iter().map(|&x| x.abs()).sum();
    let sum_im_sq: f64 = im_values.iter().map(|&x| x * x).sum();

    let numerator = sum_im * sum_im - sum_im_sq;
    let denominator = sum_abs_im * sum_abs_im - sum_im_sq;

    if denominator.abs() > 1e-15 {
        Ok((numerator / denominator).abs())
    } else {
        Ok(0.0)
    }
}

/// Compute wPLI for all channel pairs
pub fn wpli_all_pairs(
    epochs: &[Vec<Vec<f64>>],
    sfreq: f64,
    fmin: f64,
    fmax: f64,
    n_fft: usize,
) -> ConnectivityResult<Vec<Vec<f64>>> {
    if epochs.is_empty() {
        return Err(ConnectivityError::InsufficientData(
            "No epochs provided".to_string(),
        ));
    }

    let n_channels = epochs[0].len();
    let freq_indices = utils::freq_indices(sfreq, n_fft, fmin, fmax);

    if freq_indices.is_empty() {
        return Err(ConnectivityError::InvalidParameters(
            "No frequencies in range".to_string(),
        ));
    }

    let mut pairs: Vec<(usize, usize)> = Vec::new();
    for i in 0..n_channels {
        for j in (i + 1)..n_channels {
            pairs.push((i, j));
        }
    }

    let window = utils::hanning_window(epochs[0][0].len());

    // For each pair and frequency, accumulate Im values
    let n_freqs = freq_indices.len();
    let mut im_sum: Vec<Vec<f64>> = vec![vec![0.0; n_freqs]; pairs.len()];
    let mut abs_im_sum: Vec<Vec<f64>> = vec![vec![0.0; n_freqs]; pairs.len()];

    for epoch in epochs {
        // FFT all channels
        let ffts: Vec<Vec<Complex64>> = epoch
            .iter()
            .map(|ch| {
                let windowed = utils::apply_window(ch, &window);
                utils::rfft(&windowed, n_fft)
            })
            .collect();

        for (p, &(i, j)) in pairs.iter().enumerate() {
            for (f_idx, &k) in freq_indices.iter().enumerate() {
                let csd = ffts[i][k] * ffts[j][k].conj();
                let im = csd.im;
                im_sum[p][f_idx] += im;
                abs_im_sum[p][f_idx] += im.abs();
            }
        }
    }

    // Compute wPLI for each pair and frequency
    let result: Vec<Vec<f64>> = im_sum
        .iter()
        .zip(&abs_im_sum)
        .map(|(im, abs_im)| {
            im.iter()
                .zip(abs_im)
                .map(|(&s, &a)| if a > 1e-15 { s.abs() / a } else { 0.0 })
                .collect()
        })
        .collect();

    Ok(result)
}

/// Compute debiased wPLI for all channel pairs
pub fn dwpli_all_pairs(
    epochs: &[Vec<Vec<f64>>],
    sfreq: f64,
    fmin: f64,
    fmax: f64,
    n_fft: usize,
) -> ConnectivityResult<Vec<Vec<f64>>> {
    if epochs.is_empty() {
        return Err(ConnectivityError::InsufficientData(
            "No epochs provided".to_string(),
        ));
    }

    let n_channels = epochs[0].len();
    let n_epochs = epochs.len();
    let freq_indices = utils::freq_indices(sfreq, n_fft, fmin, fmax);

    if freq_indices.is_empty() {
        return Err(ConnectivityError::InvalidParameters(
            "No frequencies in range".to_string(),
        ));
    }

    let mut pairs: Vec<(usize, usize)> = Vec::new();
    for i in 0..n_channels {
        for j in (i + 1)..n_channels {
            pairs.push((i, j));
        }
    }

    let window = utils::hanning_window(epochs[0][0].len());
    let _n_freqs = freq_indices.len();

    // Collect all Im values for each pair and frequency
    let mut im_values: Vec<Vec<Vec<f64>>> = pairs
        .iter()
        .map(|_| {
            freq_indices
                .iter()
                .map(|_| Vec::with_capacity(n_epochs))
                .collect()
        })
        .collect();

    for epoch in epochs {
        let ffts: Vec<Vec<Complex64>> = epoch
            .iter()
            .map(|ch| {
                let windowed = utils::apply_window(ch, &window);
                utils::rfft(&windowed, n_fft)
            })
            .collect();

        for (p, &(i, j)) in pairs.iter().enumerate() {
            for (f_idx, &k) in freq_indices.iter().enumerate() {
                let csd = ffts[i][k] * ffts[j][k].conj();
                im_values[p][f_idx].push(csd.im);
            }
        }
    }

    // Compute dwPLI for each pair and frequency
    let result: Vec<Vec<f64>> = im_values
        .iter()
        .map(|pair_im| {
            pair_im
                .iter()
                .map(|freq_im| {
                    if freq_im.len() < 2 {
                        return 0.0;
                    }

                    let sum_im: f64 = freq_im.iter().sum();
                    let sum_abs_im: f64 = freq_im.iter().map(|&x| x.abs()).sum();
                    let sum_im_sq: f64 = freq_im.iter().map(|&x| x * x).sum();

                    let numerator = sum_im * sum_im - sum_im_sq;
                    let denominator = sum_abs_im * sum_abs_im - sum_im_sq;

                    if denominator.abs() > 1e-15 {
                        (numerator / denominator).abs()
                    } else {
                        0.0
                    }
                })
                .collect()
        })
        .collect();

    Ok(result)
}

/// Phase Lag Index (PLI)
///
/// PLI is the predecessor to wPLI, using only the sign of the phase difference
/// PLI = |⟨sign(Im(Sxy))⟩|
pub fn pli(
    x_epochs: &[Vec<f64>],
    y_epochs: &[Vec<f64>],
    sfreq: f64,
    fmin: f64,
    fmax: f64,
) -> ConnectivityResult<f64> {
    if x_epochs.len() != y_epochs.len() {
        return Err(ConnectivityError::DimensionMismatch(
            "Number of epochs must match".to_string(),
        ));
    }

    if x_epochs.is_empty() {
        return Err(ConnectivityError::InsufficientData(
            "No epochs provided".to_string(),
        ));
    }

    let n_samples = x_epochs[0].len();
    let n_fft = n_samples;
    let freq_indices = utils::freq_indices(sfreq, n_fft, fmin, fmax);

    if freq_indices.is_empty() {
        return Err(ConnectivityError::InvalidParameters(
            "No frequencies in range".to_string(),
        ));
    }

    let window = utils::hanning_window(n_samples);

    let mut sign_sum = 0.0;
    let mut count = 0;

    for (x_epoch, y_epoch) in x_epochs.iter().zip(y_epochs) {
        let x_win = utils::apply_window(x_epoch, &window);
        let y_win = utils::apply_window(y_epoch, &window);

        let x_fft = utils::rfft(&x_win, n_fft);
        let y_fft = utils::rfft(&y_win, n_fft);

        for &k in &freq_indices {
            let csd = x_fft[k] * y_fft[k].conj();
            sign_sum += csd.im.signum();
            count += 1;
        }
    }

    if count > 0 {
        Ok((sign_sum / count as f64).abs())
    } else {
        Ok(0.0)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::f64::consts::PI;

    fn create_test_epochs(
        n_epochs: usize,
        n_samples: usize,
        phase_lag: f64,
    ) -> (Vec<Vec<f64>>, Vec<Vec<f64>>) {
        let sfreq = 100.0;
        let freq = 10.0;

        let x_epochs: Vec<Vec<f64>> = (0..n_epochs)
            .map(|_| {
                (0..n_samples)
                    .map(|i| (2.0 * PI * freq * i as f64 / sfreq).sin())
                    .collect()
            })
            .collect();

        let y_epochs: Vec<Vec<f64>> = (0..n_epochs)
            .map(|_| {
                (0..n_samples)
                    .map(|i| (2.0 * PI * freq * i as f64 / sfreq + phase_lag).sin())
                    .collect()
            })
            .collect();

        (x_epochs, y_epochs)
    }

    #[test]
    fn test_wpli_zero_lag() {
        // Zero phase lag should give low wPLI (volume conduction)
        let (x, y) = create_test_epochs(20, 200, 0.0);
        let wpli_val = wpli(&x, &y, 100.0, 8.0, 12.0).unwrap();
        assert!(
            wpli_val < 0.3,
            "wPLI for zero lag should be low, got {}",
            wpli_val
        );
    }

    #[test]
    fn test_wpli_constant_lag() {
        // Constant phase lag should give high wPLI
        let (x, y) = create_test_epochs(20, 200, PI / 4.0);
        let wpli_val = wpli(&x, &y, 100.0, 8.0, 12.0).unwrap();
        assert!(
            wpli_val > 0.5,
            "wPLI for constant lag should be high, got {}",
            wpli_val
        );
    }

    #[test]
    fn test_dwpli() {
        let (x, y) = create_test_epochs(20, 200, PI / 4.0);
        let dwpli_val = dwpli(&x, &y, 100.0, 8.0, 12.0).unwrap();
        assert!(dwpli_val >= 0.0 && dwpli_val <= 1.0);
    }

    #[test]
    fn test_pli() {
        // Constant phase lead should give high PLI
        let (x, y) = create_test_epochs(20, 200, PI / 4.0);
        let pli_val = pli(&x, &y, 100.0, 8.0, 12.0).unwrap();
        assert!(
            pli_val > 0.5,
            "PLI for constant lag should be high, got {}",
            pli_val
        );
    }

    #[test]
    fn test_wpli_all_pairs() {
        let n = 200;
        let sfreq = 100.0;

        let epochs: Vec<Vec<Vec<f64>>> = (0..10)
            .map(|_| {
                (0..3)
                    .map(|ch| {
                        (0..n)
                            .map(|i| {
                                let phase = ch as f64 * PI / 8.0;
                                (2.0 * PI * 10.0 * i as f64 / sfreq + phase).sin()
                            })
                            .collect()
                    })
                    .collect()
            })
            .collect();

        let result = wpli_all_pairs(&epochs, sfreq, 8.0, 12.0, n).unwrap();

        // 3 channels = 3 pairs
        assert_eq!(result.len(), 3);
        // All wPLI values should be in [0, 1]
        assert!(
            result
                .iter()
                .all(|pair| pair.iter().all(|&v| v >= 0.0 && v <= 1.0))
        );
    }
}