//! Spectral coherence estimation.
//!
//! Coherence measures the linear relationship between two signals in the frequency domain.
//!
//! ## Mathematical Background
//!
//! Magnitude-squared coherence: Cxy(f) = |Sxy(f)|² / (Sxx(f) * Syy(f))
//!
//! where Sxy is cross-spectral density and Sxx, Syy are auto-spectral densities.

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

/// Compute coherence between two signals
///
/// # Arguments
/// * `x` - First signal [n_samples]
/// * `y` - Second signal [n_samples]
/// * `sfreq` - Sampling frequency (Hz)
/// * `n_fft` - FFT length (None = signal length)
///
/// # Returns
/// (coherence, frequencies) where coherence is [n_freqs]
pub fn coherence(
    x: &[f64],
    y: &[f64],
    sfreq: f64,
    n_fft: Option<usize>,
) -> ConnectivityResult<(Vec<f64>, Vec<f64>)> {
    if x.len() != y.len() {
        return Err(ConnectivityError::DimensionMismatch(
            "Signals must have same length".to_string(),
        ));
    }

    let n = x.len();
    let n_fft = n_fft.unwrap_or(n);

    // Apply window
    let window = utils::hanning_window(n);
    let x_win = utils::apply_window(x, &window);
    let y_win = utils::apply_window(y, &window);

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

    // Cross and auto spectral densities
    let sxy = utils::cross_spectral_density(&x_fft, &y_fft);
    let sxx = utils::power_spectral_density(&x_fft);
    let syy = utils::power_spectral_density(&y_fft);

    // Coherence: |Sxy|² / (Sxx * Syy)
    // Clamp to [0, 1] due to numerical precision
    let coh: Vec<f64> = sxy
        .iter()
        .zip(sxx.iter())
        .zip(syy.iter())
        .map(|((xy, &xx), &yy)| {
            let denom = xx * yy;
            if denom > 1e-15 {
                (xy.norm_sqr() / denom).clamp(0.0, 1.0)
            } else {
                0.0
            }
        })
        .collect();

    let freqs = utils::fft_freqs(sfreq, n_fft, 0.0, sfreq / 2.0);

    Ok((coh, freqs))
}

/// Compute coherence between specific channel pairs
///
/// # Arguments
/// * `x` - First signal [n_samples]
/// * `y` - Second signal [n_samples]
/// * `sfreq` - Sampling frequency
/// * `fmin` - Minimum frequency
/// * `fmax` - Maximum frequency
/// * `n_fft` - FFT length
pub fn coherence_pairs(
    x: &[f64],
    y: &[f64],
    sfreq: f64,
    fmin: f64,
    fmax: f64,
    n_fft: usize,
) -> ConnectivityResult<Vec<f64>> {
    let (coh, freqs) = coherence(x, y, sfreq, Some(n_fft))?;

    // Filter to frequency range
    let result: Vec<f64> = coh
        .iter()
        .zip(freqs.iter())
        .filter(|&(_, f)| *f >= fmin && *f <= fmax)
        .map(|(&c, _)| c)
        .collect();

    Ok(result)
}

/// Compute imaginary coherence (robust to volume conduction)
///
/// Imaginary coherence only considers the imaginary part of the cross-spectrum,
/// which is zero for volume-conducted signals (zero phase lag).
pub fn imaginary_coherence(
    x: &[f64],
    y: &[f64],
    sfreq: f64,
    n_fft: Option<usize>,
) -> ConnectivityResult<(Vec<f64>, Vec<f64>)> {
    if x.len() != y.len() {
        return Err(ConnectivityError::DimensionMismatch(
            "Signals must have same length".to_string(),
        ));
    }

    let n = x.len();
    let n_fft = n_fft.unwrap_or(n);

    let window = utils::hanning_window(n);
    let x_win = utils::apply_window(x, &window);
    let y_win = utils::apply_window(y, &window);

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

    let sxy = utils::cross_spectral_density(&x_fft, &y_fft);
    let sxx = utils::power_spectral_density(&x_fft);
    let syy = utils::power_spectral_density(&y_fft);

    // Imaginary coherence: Im(Sxy)² / (Sxx * Syy)
    // Clamp to [0, 1]
    let imcoh: Vec<f64> = sxy
        .iter()
        .zip(sxx.iter())
        .zip(syy.iter())
        .map(|((xy, &xx), &yy)| {
            let denom = xx * yy;
            if denom > 1e-15 {
                ((xy.im * xy.im) / denom).clamp(0.0, 1.0)
            } else {
                0.0
            }
        })
        .collect();

    let freqs = utils::fft_freqs(sfreq, n_fft, 0.0, sfreq / 2.0);

    Ok((imcoh, freqs))
}

/// Compute coherence across all channel pairs for epoched data
pub fn coherence_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 _n_freqs = freq_indices.len();

    // All channel 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 cross and auto spectra across epochs
    let mut sxy_sum: Vec<Vec<Complex64>> =
        vec![vec![Complex64::new(0.0, 0.0); n_fft / 2 + 1]; pairs.len()];
    let mut sxx_sum: Vec<Vec<f64>> = vec![vec![0.0; n_fft / 2 + 1]; n_channels];
    let _syy_sum: Vec<Vec<f64>> = vec![vec![0.0; n_fft / 2 + 1]; n_channels];

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

        // Accumulate auto-spectra
        for (ch, fft) in ffts.iter().enumerate() {
            for (i, c) in fft.iter().enumerate() {
                sxx_sum[ch][i] += c.norm_sqr();
            }
        }

        // Accumulate cross-spectra
        for (p, &(i, j)) in pairs.iter().enumerate() {
            for (k, (a, b)) in ffts[i].iter().zip(&ffts[j]).enumerate() {
                sxy_sum[p][k] += a * b.conj();
            }
        }
    }

    // Compute coherence
    let mut result = Vec::with_capacity(pairs.len());

    for (p, &(i, j)) in pairs.iter().enumerate() {
        let coh: Vec<f64> = freq_indices
            .iter()
            .map(|&k| {
                let denom = sxx_sum[i][k] * sxx_sum[j][k];
                if denom > 1e-15 {
                    (sxy_sum[p][k].norm_sqr() / denom).clamp(0.0, 1.0)
                } else {
                    0.0
                }
            })
            .collect();
        result.push(coh);
    }

    Ok(result)
}

/// Compute imaginary coherence across all channel pairs
pub fn imag_coherence_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());

    let mut sxy_sum: Vec<Vec<Complex64>> =
        vec![vec![Complex64::new(0.0, 0.0); n_fft / 2 + 1]; pairs.len()];
    let mut sxx_sum: Vec<Vec<f64>> = vec![vec![0.0; n_fft / 2 + 1]; n_channels];

    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 (ch, fft) in ffts.iter().enumerate() {
            for (i, c) in fft.iter().enumerate() {
                sxx_sum[ch][i] += c.norm_sqr();
            }
        }

        for (p, &(i, j)) in pairs.iter().enumerate() {
            for (k, (a, b)) in ffts[i].iter().zip(&ffts[j]).enumerate() {
                sxy_sum[p][k] += a * b.conj();
            }
        }
    }

    let mut result = Vec::with_capacity(pairs.len());

    for (p, &(i, j)) in pairs.iter().enumerate() {
        let imcoh: Vec<f64> = freq_indices
            .iter()
            .map(|&k| {
                let denom = sxx_sum[i][k] * sxx_sum[j][k];
                if denom > 1e-15 {
                    let im = sxy_sum[p][k].im;
                    ((im * im) / denom).clamp(0.0, 1.0)
                } else {
                    0.0
                }
            })
            .collect();
        result.push(imcoh);
    }

    Ok(result)
}

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

    fn create_test_signals() -> (Vec<f64>, Vec<f64>) {
        let n = 500;
        let sfreq = 100.0;

        // Two signals with known coherence
        let x: Vec<f64> = (0..n)
            .map(|i| {
                let t = i as f64 / sfreq;
                (2.0 * PI * 10.0 * t).sin() + 0.5 * (2.0 * PI * 20.0 * t).sin()
            })
            .collect();

        let y: Vec<f64> = (0..n)
            .map(|i| {
                let t = i as f64 / sfreq;
                0.8 * (2.0 * PI * 10.0 * t).sin() + 0.3 * (2.0 * PI * 30.0 * t).sin()
            })
            .collect();

        (x, y)
    }

    #[test]
    fn test_coherence() {
        let (x, y) = create_test_signals();

        let (coh, freqs) = coherence(&x, &y, 100.0, None).unwrap();

        // Should have positive coherence
        assert!(coh.len() > 0);
        assert!(coh.iter().all(|&c| c >= 0.0 && c <= 1.0));

        // Coherence at 10 Hz should be higher (shared component)
        let idx_10hz = freqs.iter().position(|&f| (f - 10.0).abs() < 1.0).unwrap();
        assert!(coh[idx_10hz] > 0.5);
    }

    #[test]
    fn test_perfect_coherence() {
        let n = 500;
        // Use a pure sinusoid at a specific frequency
        let x: Vec<f64> = (0..n)
            .map(|i| (2.0 * PI * 10.0 * i as f64 / 100.0).sin())
            .collect();
        let y = x.clone(); // Perfect copy

        let (coh, freqs) = coherence(&x, &y, 100.0, None).unwrap();

        // Coherence at the signal frequency should be very high
        let idx_10hz = freqs.iter().position(|&f| (f - 10.0).abs() < 1.0).unwrap();
        assert!(
            coh[idx_10hz] > 0.95,
            "Coherence at signal frequency should be ~1.0, got {}",
            coh[idx_10hz]
        );
    }

    #[test]
    fn test_imaginary_coherence() {
        let (x, y) = create_test_signals();

        let (imcoh, _) = imaginary_coherence(&x, &y, 100.0, None).unwrap();

        // Imaginary coherence should be bounded [0, 1]
        assert!(imcoh.iter().all(|&c| c >= 0.0 && c <= 1.0));
    }

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

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

        let result = coherence_all_pairs(&epochs, sfreq, 5.0, 30.0, n).unwrap();

        // 3 channels = 3 pairs
        assert_eq!(result.len(), 3);
        // All coherence values should be valid
        assert!(
            result
                .iter()
                .all(|pair| pair.iter().all(|&c| c >= 0.0 && c <= 1.0))
        );
    }
}