//! Phase Locking Value (PLV) and related phase connectivity measures.
//!
//! PLV measures the consistency of phase difference between two signals
//! across trials/epochs.
//!
//! ## Mathematical Background
//!
//! PLV = |⟨exp(i(φ_x - φ_y))⟩|
//!
//! where φ_x and φ_y are instantaneous phases, and ⟨⟩ is the average
//! across trials.

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

/// Compute PLV 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 for bandpass
/// * `fmax` - Maximum frequency for bandpass
///
/// # Returns
/// PLV value (0-1)
pub fn plv(
    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();

    // For each epoch, compute phase difference and accumulate
    let mut phase_diff_sum = vec![Complex64::new(0.0, 0.0); n_samples];

    for (x_epoch, y_epoch) in x_epochs.iter().zip(y_epochs) {
        // Bandpass filter
        let x_filt = utils::bandpass_fft(x_epoch, sfreq, fmin, fmax);
        let y_filt = utils::bandpass_fft(y_epoch, sfreq, fmin, fmax);

        // Get analytic signals (Hilbert transform)
        let x_analytic = utils::hilbert(&x_filt);
        let y_analytic = utils::hilbert(&y_filt);

        // Compute phase difference: exp(i*(φ_x - φ_y)) = x_analytic * conj(y_analytic) / |...|
        for (i, (ax, ay)) in x_analytic.iter().zip(&y_analytic).enumerate() {
            let phase_diff = ax * ay.conj();
            let norm = phase_diff.norm();
            if norm > 1e-15 {
                phase_diff_sum[i] += phase_diff / norm;
            }
        }
    }

    // PLV = |mean(exp(i*phase_diff))|
    let plv_values: Vec<f64> = phase_diff_sum
        .iter()
        .map(|c| c.norm() / n_epochs as f64)
        .collect();

    // Return mean PLV across time
    let mean_plv = plv_values.iter().sum::<f64>() / plv_values.len() as f64;
    Ok(mean_plv)
}

/// Compute PLV for epoched multi-channel data
///
/// Returns PLV time course for each channel pair
pub fn plv_epochs(
    epochs: &[Vec<Vec<f64>>], // [n_epochs][n_channels][n_samples]
    ch1: usize,
    ch2: usize,
    sfreq: f64,
    fmin: f64,
    fmax: f64,
) -> ConnectivityResult<Vec<f64>> {
    if epochs.is_empty() {
        return Err(ConnectivityError::InsufficientData(
            "No epochs provided".to_string(),
        ));
    }

    let n_channels = epochs[0].len();
    if ch1 >= n_channels || ch2 >= n_channels {
        return Err(ConnectivityError::InvalidParameters(
            "Channel index out of range".to_string(),
        ));
    }

    let n_epochs = epochs.len();
    let n_samples = epochs[0][0].len();

    let mut phase_diff_sum = vec![Complex64::new(0.0, 0.0); n_samples];

    for epoch in epochs {
        let x_filt = utils::bandpass_fft(&epoch[ch1], sfreq, fmin, fmax);
        let y_filt = utils::bandpass_fft(&epoch[ch2], sfreq, fmin, fmax);

        let x_analytic = utils::hilbert(&x_filt);
        let y_analytic = utils::hilbert(&y_filt);

        for (i, (ax, ay)) in x_analytic.iter().zip(&y_analytic).enumerate() {
            let phase_diff = ax * ay.conj();
            let norm = phase_diff.norm();
            if norm > 1e-15 {
                phase_diff_sum[i] += phase_diff / norm;
            }
        }
    }

    let plv_timecourse: Vec<f64> = phase_diff_sum
        .iter()
        .map(|c| (c.norm() / n_epochs as f64).clamp(0.0, 1.0))
        .collect();

    Ok(plv_timecourse)
}

/// Compute PLV for all channel pairs in frequency bands
pub fn plv_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);

    // Get all pairs
    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());

    // Accumulate phase difference across epochs for each frequency
    let mut phase_sum: Vec<Vec<Vec<Complex64>>> = pairs
        .iter()
        .map(|_| freq_indices.iter().map(|_| Vec::new()).collect())
        .collect();

    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 each pair and frequency, compute phase difference
        for (p, &(i, j)) in pairs.iter().enumerate() {
            for (f_idx, &k) in freq_indices.iter().enumerate() {
                let phase_diff = ffts[i][k] * ffts[j][k].conj();
                let norm = phase_diff.norm();
                if norm > 1e-15 {
                    let normalized = phase_diff / norm;
                    phase_sum[p][f_idx].push(normalized);
                }
            }
        }
    }

    // Compute PLV for each pair and frequency
    let result: Vec<Vec<f64>> = phase_sum
        .iter()
        .map(|pair_phases| {
            pair_phases
                .iter()
                .map(|freq_phases| {
                    if freq_phases.is_empty() {
                        0.0
                    } else {
                        let sum: Complex64 = freq_phases.iter().sum();
                        sum.norm() / freq_phases.len() as f64
                    }
                })
                .collect()
        })
        .collect();

    Ok(result)
}

/// Pairwise Phase Consistency (PPC)
///
/// PPC is an unbiased estimator of squared PLV
/// PPC = (n * PLV² - 1) / (n - 1)
pub fn ppc_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 n_epochs < 2 {
        return Err(ConnectivityError::InsufficientData(
            "PPC requires at least 2 epochs".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());

    // Collect all phase values
    let mut phases: Vec<Vec<Vec<Complex64>>> = pairs
        .iter()
        .map(|_| freq_indices.iter().map(|_| Vec::new()).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 phase_diff = ffts[i][k] * ffts[j][k].conj();
                let norm = phase_diff.norm();
                if norm > 1e-15 {
                    phases[p][f_idx].push(phase_diff / norm);
                }
            }
        }
    }

    // Compute PPC: unbiased estimator using pairwise products
    let result: Vec<Vec<f64>> = phases
        .iter()
        .map(|pair_phases| {
            pair_phases
                .iter()
                .map(|freq_phases| {
                    let n = freq_phases.len();
                    if n < 2 {
                        return 0.0;
                    }

                    // Sum of all pairwise products
                    let mut sum_real = 0.0;
                    for i in 0..n {
                        for j in (i + 1)..n {
                            let prod = freq_phases[i] * freq_phases[j].conj();
                            sum_real += prod.re;
                        }
                    }

                    // Number of pairs
                    let n_pairs = n * (n - 1) / 2;
                    sum_real / n_pairs as f64
                })
                .collect()
        })
        .collect();

    Ok(result)
}

/// Compute circular mean of phases
pub fn circular_mean(phases: &[f64]) -> f64 {
    if phases.is_empty() {
        return 0.0;
    }

    let sum_sin: f64 = phases.iter().map(|&p| p.sin()).sum();
    let sum_cos: f64 = phases.iter().map(|&p| p.cos()).sum();

    sum_sin.atan2(sum_cos)
}

/// Compute circular variance (related to PLV)
pub fn circular_variance(phases: &[f64]) -> f64 {
    if phases.is_empty() {
        return 0.0;
    }

    let n = phases.len() as f64;
    let sum_sin: f64 = phases.iter().map(|&p| p.sin()).sum();
    let sum_cos: f64 = phases.iter().map(|&p| p.cos()).sum();

    let r = (sum_sin * sum_sin + sum_cos * sum_cos).sqrt() / n;
    1.0 - r
}

/// Corrected Imaginary PLV (ciPLV)
///
/// ciPLV is more robust to volume conduction artifacts than standard PLV
/// by only considering the imaginary part of the cross-spectrum, then
/// normalizing to account for the loss of information.
///
/// ciPLV = |Im(PLV)| / sqrt(1 - Re(PLV)²)
///
/// This provides a measure that is:
/// - Zero for perfectly phase-locked signals with zero lag (volume conduction)
/// - Non-zero for true interactions with phase lag
pub fn ciplv(
    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 mut plv_sum = vec![Complex64::new(0.0, 0.0); n_samples];

    for (x_epoch, y_epoch) in x_epochs.iter().zip(y_epochs) {
        let x_filt = utils::bandpass_fft(x_epoch, sfreq, fmin, fmax);
        let y_filt = utils::bandpass_fft(y_epoch, sfreq, fmin, fmax);

        let x_analytic = utils::hilbert(&x_filt);
        let y_analytic = utils::hilbert(&y_filt);

        for (i, (ax, ay)) in x_analytic.iter().zip(&y_analytic).enumerate() {
            let phase_diff = ax * ay.conj();
            let norm = phase_diff.norm();
            if norm > 1e-15 {
                plv_sum[i] += phase_diff / norm;
            }
        }
    }

    // Compute ciPLV from complex PLV
    let ciplv_values: Vec<f64> = plv_sum
        .iter()
        .map(|c| {
            let plv_complex = *c / n_epochs as f64;
            let re = plv_complex.re;
            let im = plv_complex.im;

            // ciPLV = |Im(PLV)| / sqrt(1 - Re(PLV)²)
            let denom = (1.0 - re * re).sqrt();
            if denom > 1e-10 { im.abs() / denom } else { 0.0 }
        })
        .collect();

    let mean_ciplv = ciplv_values.iter().sum::<f64>() / ciplv_values.len() as f64;
    Ok(mean_ciplv.clamp(0.0, 1.0))
}

/// Weighted PLV (wPLV)
///
/// wPLV weights each epoch's contribution by the amplitude of the signals,
/// giving more weight to epochs with stronger oscillations.
///
/// wPLV = |Σ(|A_x||A_y| * exp(i(φ_x - φ_y)))| / Σ(|A_x||A_y|)
///
/// This is useful when:
/// - Signal quality varies across epochs
/// - Oscillation amplitude is informative
/// - You want to down-weight noisy epochs
pub fn wplv(
    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 mut weighted_sum = vec![Complex64::new(0.0, 0.0); n_samples];
    let mut weight_sum = vec![0.0; n_samples];

    for (x_epoch, y_epoch) in x_epochs.iter().zip(y_epochs) {
        let x_filt = utils::bandpass_fft(x_epoch, sfreq, fmin, fmax);
        let y_filt = utils::bandpass_fft(y_epoch, sfreq, fmin, fmax);

        let x_analytic = utils::hilbert(&x_filt);
        let y_analytic = utils::hilbert(&y_filt);

        for (i, (ax, ay)) in x_analytic.iter().zip(&y_analytic).enumerate() {
            let amp_x = ax.norm();
            let amp_y = ay.norm();
            let weight = amp_x * amp_y;

            if weight > 1e-15 {
                // Phase difference
                let phase_diff = ax * ay.conj();
                let phase_unit = phase_diff / phase_diff.norm();

                weighted_sum[i] += phase_unit * weight;
                weight_sum[i] += weight;
            }
        }
    }

    // Compute weighted PLV
    let wplv_values: Vec<f64> = weighted_sum
        .iter()
        .zip(&weight_sum)
        .map(|(ws, &wt)| if wt > 1e-15 { ws.norm() / wt } else { 0.0 })
        .collect();

    let mean_wplv = wplv_values.iter().sum::<f64>() / wplv_values.len() as f64;
    Ok(mean_wplv.clamp(0.0, 1.0))
}

/// Compute ciPLV for all channel pairs
pub fn ciplv_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);

    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 mut plv_sum: Vec<Vec<Vec<Complex64>>> = pairs
        .iter()
        .map(|_| freq_indices.iter().map(|_| Vec::new()).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 phase_diff = ffts[i][k] * ffts[j][k].conj();
                let norm = phase_diff.norm();
                if norm > 1e-15 {
                    let normalized = phase_diff / norm;
                    plv_sum[p][f_idx].push(normalized);
                }
            }
        }
    }

    // Compute ciPLV
    let result: Vec<Vec<f64>> = plv_sum
        .iter()
        .map(|pair_phases| {
            pair_phases
                .iter()
                .map(|freq_phases| {
                    if freq_phases.is_empty() {
                        return 0.0;
                    }

                    let sum: Complex64 = freq_phases.iter().sum();
                    let plv_complex = sum / freq_phases.len() as f64;

                    let re = plv_complex.re;
                    let im = plv_complex.im;

                    let denom = (1.0 - re * re).sqrt();
                    if denom > 1e-10 {
                        (im.abs() / denom).clamp(0.0, 1.0)
                    } else {
                        0.0
                    }
                })
                .collect()
        })
        .collect();

    Ok(result)
}

/// Compute wPLV for all channel pairs
pub fn wplv_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);

    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());

    // Accumulate weighted sums
    let mut weighted_sums: Vec<Vec<Complex64>> = pairs
        .iter()
        .map(|_| vec![Complex64::new(0.0, 0.0); freq_indices.len()])
        .collect();
    let mut weight_sums: Vec<Vec<f64>> = pairs
        .iter()
        .map(|_| vec![0.0; freq_indices.len()])
        .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 amp_i = ffts[i][k].norm();
                let amp_j = ffts[j][k].norm();
                let weight = amp_i * amp_j;

                if weight > 1e-15 {
                    let phase_diff = ffts[i][k] * ffts[j][k].conj();
                    let phase_unit = phase_diff / phase_diff.norm();

                    weighted_sums[p][f_idx] += phase_unit * weight;
                    weight_sums[p][f_idx] += weight;
                }
            }
        }
    }

    // Compute wPLV
    let result: Vec<Vec<f64>> = weighted_sums
        .iter()
        .zip(&weight_sums)
        .map(|(ws_pair, wt_pair)| {
            ws_pair
                .iter()
                .zip(wt_pair)
                .map(|(ws, &wt)| {
                    if wt > 1e-15 {
                        (ws.norm() / wt).clamp(0.0, 1.0)
                    } else {
                        0.0
                    }
                })
                .collect()
        })
        .collect();

    Ok(result)
}

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

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

        // Create perfectly synchronized signals
        let epochs: Vec<Vec<f64>> = (0..10)
            .map(|_| {
                (0..n)
                    .map(|i| (2.0 * PI * 10.0 * i as f64 / sfreq).sin())
                    .collect()
            })
            .collect();

        // Same signal = perfect PLV
        let plv_val = plv(&epochs, &epochs, sfreq, 8.0, 12.0).unwrap();
        assert!(
            plv_val > 0.9,
            "PLV for identical signals should be ~1.0, got {}",
            plv_val
        );
    }

    #[test]
    fn test_plv_random() {
        let n = 200;

        // Random signals should have low PLV
        let x_epochs: Vec<Vec<f64>> = (0..20)
            .map(|seed| {
                (0..n)
                    .map(|i| ((i as f64 * 0.1 + seed as f64).sin() * 1000.0).sin())
                    .collect()
            })
            .collect();

        let y_epochs: Vec<Vec<f64>> = (0..20)
            .map(|seed| {
                (0..n)
                    .map(|i| ((i as f64 * 0.2 + seed as f64 * 2.0).cos() * 1000.0).cos())
                    .collect()
            })
            .collect();

        let plv_val = plv(&x_epochs, &y_epochs, 100.0, 8.0, 12.0).unwrap();
        assert!(
            plv_val < 0.5,
            "PLV for random signals should be low, got {}",
            plv_val
        );
    }

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

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

        let plv_tc = plv_epochs(&epochs, 0, 1, sfreq, 8.0, 12.0).unwrap();
        assert_eq!(plv_tc.len(), n);
        assert!(plv_tc.iter().all(|&v| v >= 0.0 && v <= 1.0));
    }

    #[test]
    fn test_circular_mean() {
        let phases = vec![0.0, PI / 2.0, 0.0];
        let mean = circular_mean(&phases);
        assert!((mean - PI / 6.0).abs() < 0.5);
    }

    #[test]
    fn test_circular_variance() {
        // All same phase = zero variance
        let phases = vec![0.0, 0.0, 0.0];
        let var = circular_variance(&phases);
        assert!(var.abs() < 1e-10);

        // Opposite phases = high variance
        let phases2 = vec![0.0, PI, 0.0, PI];
        let var2 = circular_variance(&phases2);
        assert!(var2 > 0.5);
    }

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

        // Create phase-shifted signals (non-zero phase lag)
        let epochs: Vec<Vec<f64>> = (0..10)
            .map(|_| {
                (0..n)
                    .map(|i| (2.0 * PI * 10.0 * i as f64 / sfreq).sin())
                    .collect()
            })
            .collect();

        // Shifted version (90 degree phase lag)
        let epochs_shifted: Vec<Vec<f64>> = (0..10)
            .map(|_| {
                (0..n)
                    .map(|i| (2.0 * PI * 10.0 * i as f64 / sfreq + PI / 2.0).sin())
                    .collect()
            })
            .collect();

        // ciPLV should be non-zero for phase-lagged signals
        let ci = ciplv(&epochs, &epochs_shifted, sfreq, 8.0, 12.0).unwrap();
        assert!(ci >= 0.0 && ci <= 1.0);
    }

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

        // Create synchronized signals
        let epochs: Vec<Vec<f64>> = (0..10)
            .map(|_| {
                (0..n)
                    .map(|i| (2.0 * PI * 10.0 * i as f64 / sfreq).sin())
                    .collect()
            })
            .collect();

        // Same signal = high wPLV
        let w = wplv(&epochs, &epochs, sfreq, 8.0, 12.0).unwrap();
        assert!(
            w > 0.8,
            "wPLV for identical signals should be high, got {}",
            w
        );
    }
}