//! Time-Frequency Representations (TFR)
//!
//! This module provides time-frequency analysis methods for MEG/EEG data:
//!
//! - **Morlet Wavelets**: Complex Morlet wavelet transform
//! - **STFT**: Short-Time Fourier Transform with configurable windows
//! - **PSD**: Power Spectral Density (Welch, Multitaper)
//!
//! # Example
//!
//! ```ignore
//! use rtx_neuro_signal::tfr::{tfr_morlet, TfrOutput};
//!
//! let data = vec![0.0; 1000];  // [n_channels x n_samples]
//! let sfreq = 256.0;
//! let freqs = vec![8.0, 10.0, 12.0, 14.0];  // Alpha band
//!
//! let tfr = tfr_morlet(&[data], sfreq, &freqs, 7.0, TfrOutput::Power)?;
//! ```

use num_complex::Complex64;
use rayon::prelude::*;
use rustfft::{Fft, FftPlanner};
use std::f64::consts::PI;
use std::sync::Arc;

use crate::{SignalError, SignalResult};

/// Output type for time-frequency representations
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum TfrOutput {
    /// Power (magnitude squared)
    Power,
    /// Phase angle (radians)
    Phase,
    /// Complex values
    Complex,
    /// Inter-trial coherence (requires multiple trials)
    Itc,
}

/// Cycle specification for wavelets
#[derive(Debug, Clone)]
pub enum CycleSpec {
    /// Fixed number of cycles for all frequencies
    Fixed(f64),
    /// Frequency-adaptive cycles (min at low freq, max at high freq)
    Adaptive {
        /// Minimum cycles at lowest frequency
        min: f64,
        /// Maximum cycles at highest frequency
        max: f64,
    },
}

impl CycleSpec {
    /// Get number of cycles for a given frequency
    pub fn get_cycles(&self, freq: f64, freq_min: f64, freq_max: f64) -> f64 {
        match self {
            CycleSpec::Fixed(n) => *n,
            CycleSpec::Adaptive { min, max } => {
                let t = (freq - freq_min) / (freq_max - freq_min + 1e-10);
                min + t * (max - min)
            }
        }
    }
}

/// Result of time-frequency analysis
#[derive(Debug, Clone)]
pub struct TfrResult {
    /// TFR data: [n_channels x n_freqs x n_times] for single trial
    /// or [n_epochs x n_channels x n_freqs x n_times] for epochs
    pub data: Vec<Vec<Vec<f64>>>,
    /// Frequencies analyzed
    pub freqs: Vec<f64>,
    /// Time points (relative to epoch start)
    pub times: Vec<f64>,
    /// Sampling frequency
    pub sfreq: f64,
    /// Output type
    pub output: TfrOutput,
}

/// Generate a complex Morlet wavelet
///
/// The Morlet wavelet is defined as:
/// w(t) = A * exp(-t²/(2σ²)) * exp(2πif₀t)
///
/// where σ = n_cycles / (2πf₀)
fn morlet_wavelet(freq: f64, sfreq: f64, n_cycles: f64) -> Vec<Complex64> {
    // Gaussian standard deviation in samples
    let sigma_t = n_cycles / (2.0 * PI * freq);
    let sigma_samples = sigma_t * sfreq;

    // Wavelet length: 5 sigma on each side
    let half_len = (5.0 * sigma_samples).ceil() as usize;
    let len = 2 * half_len + 1;

    let mut wavelet = Vec::with_capacity(len);
    let norm = 1.0 / (sigma_t * (2.0 * PI).sqrt()).sqrt();

    for i in 0..len {
        let t = (i as f64 - half_len as f64) / sfreq;
        let gaussian = (-t * t / (2.0 * sigma_t * sigma_t)).exp();
        let oscillation = Complex64::new(0.0, 2.0 * PI * freq * t).exp();
        wavelet.push(norm * gaussian * oscillation);
    }

    wavelet
}

/// Compute Morlet wavelet time-frequency representation
///
/// # Arguments
///
/// * `data` - Input data [n_channels][n_samples]
/// * `sfreq` - Sampling frequency in Hz
/// * `freqs` - Frequencies to analyze
/// * `n_cycles` - Number of wavelet cycles (typically 7)
/// * `output` - Output type (Power, Phase, Complex)
///
/// # Returns
///
/// TfrResult with data [n_channels x n_freqs x n_times]
pub fn tfr_morlet(
    data: &[Vec<f64>],
    sfreq: f64,
    freqs: &[f64],
    n_cycles: f64,
    output: TfrOutput,
) -> SignalResult<TfrResult> {
    tfr_morlet_adaptive(data, sfreq, freqs, CycleSpec::Fixed(n_cycles), output, 1)
}

/// Compute Morlet wavelet TFR with adaptive cycles and decimation
pub fn tfr_morlet_adaptive(
    data: &[Vec<f64>],
    sfreq: f64,
    freqs: &[f64],
    cycles: CycleSpec,
    output: TfrOutput,
    decim: usize,
) -> SignalResult<TfrResult> {
    if data.is_empty() || data[0].is_empty() {
        return Err(SignalError::InvalidLength("Empty data".to_string()));
    }
    if freqs.is_empty() {
        return Err(SignalError::InvalidParameters(
            "No frequencies specified".to_string(),
        ));
    }

    let n_samples = data[0].len();
    let n_freqs = freqs.len();
    let freq_min = freqs.iter().copied().fold(f64::INFINITY, f64::min);
    let freq_max = freqs.iter().copied().fold(f64::NEG_INFINITY, f64::max);

    // Decimated output length
    let n_times = n_samples.div_ceil(decim);

    // Generate wavelets for each frequency
    let wavelets: Vec<Vec<Complex64>> = freqs
        .iter()
        .map(|&f| {
            let nc = cycles.get_cycles(f, freq_min, freq_max);
            morlet_wavelet(f, sfreq, nc)
        })
        .collect();

    // Process each channel in parallel
    let result: Vec<Vec<Vec<f64>>> = data
        .par_iter()
        .map(|channel| {
            let mut freq_results = Vec::with_capacity(n_freqs);

            for wavelet in &wavelets {
                // Convolve channel with wavelet
                let convolved = convolve_complex(channel, wavelet);

                // Decimate and compute output
                let mut times_result = Vec::with_capacity(n_times);
                for t in (0..n_samples).step_by(decim) {
                    let c = convolved[t + wavelet.len() / 2];
                    let val = match output {
                        TfrOutput::Power => c.norm_sqr(),
                        TfrOutput::Phase => c.arg(),
                        TfrOutput::Complex => c.re, // Store real part; use separate fn for complex
                        TfrOutput::Itc => c.norm(), // For single trial, just amplitude
                    };
                    times_result.push(val);
                }
                freq_results.push(times_result);
            }

            freq_results
        })
        .collect();

    // Generate time vector
    let times: Vec<f64> = (0..n_times).map(|t| (t * decim) as f64 / sfreq).collect();

    Ok(TfrResult {
        data: result,
        freqs: freqs.to_vec(),
        times,
        sfreq,
        output,
    })
}

/// Convolve signal with complex wavelet
fn convolve_complex(signal: &[f64], wavelet: &[Complex64]) -> Vec<Complex64> {
    let n = signal.len();
    let w_len = wavelet.len();
    let out_len = n + w_len - 1;

    // Use FFT for efficient convolution
    let fft_len = out_len.next_power_of_two();

    let mut planner = FftPlanner::new();
    let fft = planner.plan_fft_forward(fft_len);
    let ifft = planner.plan_fft_inverse(fft_len);

    // Zero-pad signal
    let mut signal_fft: Vec<Complex64> = signal
        .iter()
        .map(|&x| Complex64::new(x, 0.0))
        .chain(std::iter::repeat(Complex64::new(0.0, 0.0)))
        .take(fft_len)
        .collect();

    // Zero-pad wavelet
    let mut wavelet_fft: Vec<Complex64> = wavelet
        .iter()
        .copied()
        .chain(std::iter::repeat(Complex64::new(0.0, 0.0)))
        .take(fft_len)
        .collect();

    // Forward FFT
    fft.process(&mut signal_fft);
    fft.process(&mut wavelet_fft);

    // Multiply in frequency domain
    for i in 0..fft_len {
        signal_fft[i] *= wavelet_fft[i];
    }

    // Inverse FFT
    ifft.process(&mut signal_fft);

    // Normalize and take valid portion
    let scale = 1.0 / fft_len as f64;
    signal_fft
        .into_iter()
        .take(out_len)
        .map(|c| c * scale)
        .collect()
}

/// Window function types for STFT
#[derive(Debug, Clone, Copy)]
pub enum Window {
    /// Rectangular (no windowing)
    Rectangular,
    /// Hann window
    Hann,
    /// Hamming window
    Hamming,
    /// Blackman window
    Blackman,
}

impl Window {
    /// Generate window function of given length
    pub fn generate(&self, len: usize) -> Vec<f64> {
        match self {
            Window::Rectangular => vec![1.0; len],
            Window::Hann => (0..len)
                .map(|i| 0.5 * (1.0 - (2.0 * PI * i as f64 / (len - 1) as f64).cos()))
                .collect(),
            Window::Hamming => (0..len)
                .map(|i| 0.54 - 0.46 * (2.0 * PI * i as f64 / (len - 1) as f64).cos())
                .collect(),
            Window::Blackman => (0..len)
                .map(|i| {
                    let a0 = 0.42;
                    let a1 = 0.5;
                    let a2 = 0.08;
                    let t = 2.0 * PI * i as f64 / (len - 1) as f64;
                    a0 - a1 * t.cos() + a2 * (2.0 * t).cos()
                })
                .collect(),
        }
    }
}

/// Result of STFT computation
#[derive(Debug, Clone)]
pub struct StftResult {
    /// Complex STFT data: [n_channels x n_freqs x n_frames]
    pub data: Vec<Vec<Vec<Complex64>>>,
    /// Frequency bins
    pub freqs: Vec<f64>,
    /// Time points (frame centers)
    pub times: Vec<f64>,
    /// Sampling frequency
    pub sfreq: f64,
}

impl StftResult {
    /// Get power spectrogram (magnitude squared)
    pub fn power(&self) -> Vec<Vec<Vec<f64>>> {
        self.data
            .iter()
            .map(|ch| {
                ch.iter()
                    .map(|f| f.iter().map(num_complex::Complex::norm_sqr).collect())
                    .collect()
            })
            .collect()
    }

    /// Get phase spectrogram
    pub fn phase(&self) -> Vec<Vec<Vec<f64>>> {
        self.data
            .iter()
            .map(|ch| {
                ch.iter()
                    .map(|f| f.iter().map(|c| c.arg()).collect())
                    .collect()
            })
            .collect()
    }
}

/// Compute Short-Time Fourier Transform
///
/// # Arguments
///
/// * `data` - Input data [n_channels][n_samples]
/// * `sfreq` - Sampling frequency
/// * `n_fft` - FFT size
/// * `hop_length` - Hop length between frames (default: n_fft / 4)
/// * `window` - Window function
///
/// # Returns
///
/// StftResult with complex STFT data
pub fn stft(
    data: &[Vec<f64>],
    sfreq: f64,
    n_fft: usize,
    hop_length: Option<usize>,
    window: Window,
) -> SignalResult<StftResult> {
    if data.is_empty() || data[0].is_empty() {
        return Err(SignalError::InvalidLength("Empty data".to_string()));
    }
    if n_fft == 0 || !n_fft.is_power_of_two() {
        return Err(SignalError::InvalidParameters(
            "n_fft must be a power of 2".to_string(),
        ));
    }

    let n_samples = data[0].len();
    let hop = hop_length.unwrap_or(n_fft / 4);
    let win = window.generate(n_fft);

    // Number of frames
    let n_frames = (n_samples.saturating_sub(n_fft)) / hop + 1;
    let n_freqs = n_fft / 2 + 1; // One-sided spectrum

    let mut planner = FftPlanner::new();
    let fft: Arc<dyn Fft<f64>> = planner.plan_fft_forward(n_fft);

    // Process each channel in parallel
    let result: Vec<Vec<Vec<Complex64>>> = data
        .par_iter()
        .map(|channel| {
            let mut freq_data = vec![vec![Complex64::new(0.0, 0.0); n_frames]; n_freqs];

            for frame_idx in 0..n_frames {
                let start = frame_idx * hop;

                // Apply window
                let mut frame: Vec<Complex64> = (0..n_fft)
                    .map(|i| {
                        let sample = if start + i < channel.len() {
                            channel[start + i]
                        } else {
                            0.0
                        };
                        Complex64::new(sample * win[i], 0.0)
                    })
                    .collect();

                // FFT
                fft.process(&mut frame);

                // Store one-sided spectrum
                for (f_idx, freq_bin) in frame.iter().take(n_freqs).enumerate() {
                    freq_data[f_idx][frame_idx] = *freq_bin;
                }
            }

            freq_data
        })
        .collect();

    // Generate frequency and time vectors
    let freqs: Vec<f64> = (0..n_freqs)
        .map(|i| i as f64 * sfreq / n_fft as f64)
        .collect();
    let times: Vec<f64> = (0..n_frames)
        .map(|i| (i * hop + n_fft / 2) as f64 / sfreq)
        .collect();

    Ok(StftResult {
        data: result,
        freqs,
        times,
        sfreq,
    })
}

/// Result of PSD computation
#[derive(Debug, Clone)]
pub struct PsdResult {
    /// Power spectral density [n_channels x n_freqs]
    pub psd: Vec<Vec<f64>>,
    /// Frequency bins
    pub freqs: Vec<f64>,
}

/// Compute Power Spectral Density using Welch's method
///
/// # Arguments
///
/// * `data` - Input data [n_channels][n_samples]
/// * `sfreq` - Sampling frequency
/// * `n_fft` - FFT size (default: 256)
/// * `n_overlap` - Overlap between segments (default: n_fft / 2)
/// * `window` - Window function
///
/// # Returns
///
/// PsdResult with PSD values
pub fn psd_welch(
    data: &[Vec<f64>],
    sfreq: f64,
    n_fft: Option<usize>,
    n_overlap: Option<usize>,
    window: Window,
) -> SignalResult<PsdResult> {
    let n_fft = n_fft.unwrap_or(256);
    let n_overlap = n_overlap.unwrap_or(n_fft / 2);
    let hop = n_fft - n_overlap;

    // Compute STFT
    let stft_result = stft(data, sfreq, n_fft, Some(hop), window)?;

    // Average power across time
    let psd: Vec<Vec<f64>> = stft_result
        .data
        .iter()
        .map(|ch| {
            ch.iter()
                .map(|freq_frames| {
                    let n_frames = freq_frames.len();
                    if n_frames == 0 {
                        return 0.0;
                    }
                    let sum: f64 = freq_frames.iter().map(num_complex::Complex::norm_sqr).sum();
                    sum / n_frames as f64
                })
                .collect()
        })
        .collect();

    Ok(PsdResult {
        psd,
        freqs: stft_result.freqs,
    })
}

/// Generate DPSS (Discrete Prolate Spheroidal Sequences) tapers
///
/// Simplified implementation using Slepian sequences approximation
fn dpss_tapers(n: usize, _bandwidth: f64, n_tapers: usize) -> Vec<Vec<f64>> {
    // Simplified: use sine tapers as approximation
    // Full DPSS would require eigenvalue computation
    let mut tapers = Vec::with_capacity(n_tapers);

    for k in 0..n_tapers {
        let taper: Vec<f64> = (0..n)
            .map(|i| {
                let t = (i as f64 + 0.5) / n as f64;
                ((k + 1) as f64 * PI * t).sin() * (2.0 / n as f64).sqrt()
            })
            .collect();
        tapers.push(taper);
    }

    tapers
}

/// Compute Power Spectral Density using multitaper method
///
/// # Arguments
///
/// * `data` - Input data [n_channels][n_samples]
/// * `sfreq` - Sampling frequency
/// * `bandwidth` - Frequency bandwidth (Hz)
/// * `fmin` - Minimum frequency (optional)
/// * `fmax` - Maximum frequency (optional)
///
/// # Returns
///
/// PsdResult with PSD values
pub fn psd_multitaper(
    data: &[Vec<f64>],
    sfreq: f64,
    bandwidth: f64,
    fmin: Option<f64>,
    fmax: Option<f64>,
) -> SignalResult<PsdResult> {
    if data.is_empty() || data[0].is_empty() {
        return Err(SignalError::InvalidLength("Empty data".to_string()));
    }

    let n_samples = data[0].len();
    let n_fft = n_samples.next_power_of_two();
    let n_freqs = n_fft / 2 + 1;

    // Calculate number of tapers
    let nw = bandwidth * n_samples as f64 / sfreq;
    let n_tapers = ((2.0 * nw - 1.0).floor() as usize).max(1);

    // Generate tapers
    let tapers = dpss_tapers(n_samples, bandwidth, n_tapers);

    let mut planner = FftPlanner::new();
    let fft: Arc<dyn Fft<f64>> = planner.plan_fft_forward(n_fft);

    // Process each channel
    let psd: Vec<Vec<f64>> = data
        .par_iter()
        .map(|channel| {
            let mut spectrum = vec![0.0; n_freqs];

            // Average across tapers
            for taper in &tapers {
                // Apply taper
                let mut tapered: Vec<Complex64> = channel
                    .iter()
                    .zip(taper.iter())
                    .map(|(&x, &t)| Complex64::new(x * t, 0.0))
                    .chain(std::iter::repeat(Complex64::new(0.0, 0.0)))
                    .take(n_fft)
                    .collect();

                // FFT
                fft.process(&mut tapered);

                // Accumulate power
                for (i, c) in tapered.iter().take(n_freqs).enumerate() {
                    spectrum[i] += c.norm_sqr();
                }
            }

            // Average and normalize
            for val in &mut spectrum {
                *val /= n_tapers as f64;
                *val /= sfreq; // Normalize to PSD units
            }

            spectrum
        })
        .collect();

    // Generate frequencies
    let freqs: Vec<f64> = (0..n_freqs)
        .map(|i| i as f64 * sfreq / n_fft as f64)
        .collect();

    // Apply frequency limits
    let fmin = fmin.unwrap_or(0.0);
    let fmax = fmax.unwrap_or(sfreq / 2.0);

    let (psd_filtered, freqs_filtered): (Vec<Vec<f64>>, Vec<f64>) = {
        let mask: Vec<bool> = freqs.iter().map(|&f| f >= fmin && f <= fmax).collect();

        let freqs_f: Vec<f64> = freqs
            .iter()
            .zip(&mask)
            .filter_map(|(&f, &m)| if m { Some(f) } else { None })
            .collect();

        let psd_f: Vec<Vec<f64>> = psd
            .iter()
            .map(|ch| {
                ch.iter()
                    .zip(&mask)
                    .filter_map(|(&p, &m)| if m { Some(p) } else { None })
                    .collect()
            })
            .collect();

        (psd_f, freqs_f)
    };

    Ok(PsdResult {
        psd: psd_filtered,
        freqs: freqs_filtered,
    })
}

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

    fn generate_sine(freq: f64, sfreq: f64, duration: f64) -> Vec<f64> {
        let n_samples = (duration * sfreq) as usize;
        (0..n_samples)
            .map(|i| (2.0 * PI * freq * i as f64 / sfreq).sin())
            .collect()
    }

    #[test]
    fn test_morlet_wavelet() {
        let wavelet = morlet_wavelet(10.0, 256.0, 7.0);
        assert!(!wavelet.is_empty());
        // Wavelet should be roughly symmetric
        let mid = wavelet.len() / 2;
        assert_relative_eq!(wavelet[mid].norm(), wavelet[mid].norm(), epsilon = 1e-10);
    }

    #[test]
    fn test_tfr_morlet_sine() {
        let sfreq = 256.0;
        let signal = generate_sine(10.0, sfreq, 2.0);
        let data = vec![signal];
        let freqs = vec![5.0, 10.0, 15.0, 20.0];

        let result = tfr_morlet(&data, sfreq, &freqs, 7.0, TfrOutput::Power).unwrap();

        assert_eq!(result.data.len(), 1); // 1 channel
        assert_eq!(result.data[0].len(), 4); // 4 frequencies
        assert_eq!(result.freqs.len(), 4);

        // 10 Hz should have highest power
        let powers: Vec<f64> = result.data[0]
            .iter()
            .map(|f| f.iter().sum::<f64>() / f.len() as f64)
            .collect();

        let max_idx = powers
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
            .map(|(i, _)| i)
            .unwrap();

        assert_eq!(max_idx, 1); // Index 1 = 10 Hz
    }

    #[test]
    fn test_stft() {
        let sfreq = 256.0;
        let signal = generate_sine(20.0, sfreq, 1.0);
        let data = vec![signal];

        let result = stft(&data, sfreq, 64, None, Window::Hann).unwrap();

        assert_eq!(result.data.len(), 1);
        assert!(!result.freqs.is_empty());
        assert!(!result.times.is_empty());

        // Check power has peak at 20 Hz
        let power = result.power();
        let mean_power: Vec<f64> = power[0]
            .iter()
            .map(|f| f.iter().sum::<f64>() / f.len() as f64)
            .collect();

        let max_idx = mean_power
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
            .map(|(i, _)| i)
            .unwrap();

        let peak_freq = result.freqs[max_idx];
        assert!(peak_freq >= 18.0 && peak_freq <= 22.0);
    }

    #[test]
    fn test_psd_welch() {
        let sfreq = 256.0;
        let signal = generate_sine(30.0, sfreq, 2.0);
        let data = vec![signal];

        let result = psd_welch(&data, sfreq, Some(128), None, Window::Hann).unwrap();

        assert_eq!(result.psd.len(), 1);
        assert!(!result.freqs.is_empty());

        // Find peak frequency
        let max_idx = result.psd[0]
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
            .map(|(i, _)| i)
            .unwrap();

        let peak_freq = result.freqs[max_idx];
        assert!(peak_freq >= 28.0 && peak_freq <= 32.0);
    }

    #[test]
    fn test_psd_multitaper() {
        let sfreq = 256.0;
        let signal = generate_sine(15.0, sfreq, 2.0);
        let data = vec![signal];

        let result = psd_multitaper(&data, sfreq, 4.0, Some(1.0), Some(50.0)).unwrap();

        assert_eq!(result.psd.len(), 1);
        assert!(!result.freqs.is_empty());

        // All frequencies should be in range
        assert!(result.freqs.iter().all(|&f| f >= 1.0 && f <= 50.0));
    }

    #[test]
    fn test_window_functions() {
        let len = 64;

        let hann = Window::Hann.generate(len);
        assert_eq!(hann.len(), len);
        assert_relative_eq!(hann[0], 0.0, epsilon = 1e-10);
        // Hann window maximum is near 1.0 but not exactly due to discrete sampling
        assert!(hann[len / 2] > 0.99);

        let hamming = Window::Hamming.generate(len);
        assert_eq!(hamming.len(), len);
        assert!(hamming[0] > 0.0); // Hamming doesn't go to 0

        let blackman = Window::Blackman.generate(len);
        assert_eq!(blackman.len(), len);
    }
}