//! Iterative Amplitude Adjusted Fourier Transform (IAAFT) surrogate generation.
//!
//! Implements the algorithm from Schreiber & Schmitz (1996/2000).
//! IAAFT surrogates preserve the amplitude spectrum (power spectrum) AND
//! marginal distribution of the original series while destroying nonlinear
//! temporal correlations. This is the gold standard for testing whether
//! an estimator detects genuine nonlinear structure vs linear autocorrelation.
//!
//! Requires `rustfft` for FFT/IFFT operations.

use rustfft::{FftPlanner, num_complex::Complex};

/// Generate an IAAFT surrogate of the input time series.
///
/// The surrogate preserves the original's amplitude spectrum and value distribution
/// while randomizing phases (destroying nonlinear correlations).
///
/// # Arguments
/// * `x` - Input time series
/// * `seed` - Random seed for initial phase randomization
/// * `max_iter` - Maximum iterations (typically 100-500; convergence is usually fast)
///
/// # Returns
/// IAAFT surrogate series of the same length
pub fn iaaft_surrogate(x: &[f64], seed: u64, max_iter: usize) -> Vec<f64> {
    let n = x.len();
    if n < 4 {
        return x.to_vec();
    }

    // Step 1: Compute amplitude spectrum of original
    let mut planner = FftPlanner::new();
    let fft = planner.plan_fft_forward(n);
    let ifft = planner.plan_fft_inverse(n);

    let mut original_fft: Vec<Complex<f64>> = x.iter().map(|&v| Complex::new(v, 0.0)).collect();
    fft.process(&mut original_fft);
    let amplitudes: Vec<f64> = original_fft.iter().map(|c| c.norm()).collect();

    // Step 2: Sort original values (for rank-ordering)
    let mut sorted_x: Vec<f64> = x.to_vec();
    sorted_x.sort_by(|a, b| a.partial_cmp(b).unwrap());

    // Step 3: Initialize with random permutation of x
    let mut surrogate = random_permutation(x, seed);

    // Step 4: Iterate
    let inv_n = 1.0 / n as f64;

    for _ in 0..max_iter {
        // 4a: FFT surrogate, replace amplitudes with original's, keep phases
        let mut s_fft: Vec<Complex<f64>> =
            surrogate.iter().map(|&v| Complex::new(v, 0.0)).collect();
        fft.process(&mut s_fft);

        for (coeff, &amp) in s_fft.iter_mut().zip(amplitudes.iter()) {
            let phase = coeff.arg();
            *coeff = Complex::from_polar(amp, phase);
        }

        // IFFT to get amplitude-adjusted series
        ifft.process(&mut s_fft);
        let spectral: Vec<f64> = s_fft.iter().map(|c| c.re * inv_n).collect();

        // 4b: Rank-order to match original distribution
        let prev_surrogate = surrogate.clone();
        surrogate = rank_order(&spectral, &sorted_x);

        // 4c: Check convergence (rank ordering unchanged)
        if surrogate == prev_surrogate {
            break;
        }
    }

    surrogate
}

/// Random permutation using Fisher-Yates shuffle with LCG
fn random_permutation(x: &[f64], seed: u64) -> Vec<f64> {
    let mut result = x.to_vec();
    let mut state = seed;

    for i in (1..result.len()).rev() {
        state = state
            .wrapping_mul(6_364_136_223_846_793_005)
            .wrapping_add(1_442_695_040_888_963_407);
        let j = (state >> 33) as usize % (i + 1);
        result.swap(i, j);
    }

    result
}

/// Rank-order `values` to match the distribution of `sorted_target`.
/// Replaces the i-th smallest value in `values` with `sorted_target[i]`.
fn rank_order(values: &[f64], sorted_target: &[f64]) -> Vec<f64> {
    let n = values.len();

    // Get rank indices of values
    let mut indices: Vec<usize> = (0..n).collect();
    indices.sort_by(|&a, &b| values[a].partial_cmp(&values[b]).unwrap());

    // Assign sorted target values according to rank
    let mut result = vec![0.0; n];
    for (rank, &idx) in indices.iter().enumerate() {
        result[idx] = sorted_target[rank];
    }

    result
}

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

    #[test]
    fn test_iaaft_preserves_length() {
        let x: Vec<f64> = (0..128).map(|i| (i as f64 * 0.1).sin()).collect();
        let s = iaaft_surrogate(&x, 42, 100);
        assert_eq!(s.len(), x.len());
    }

    #[test]
    fn test_iaaft_preserves_distribution() {
        let x: Vec<f64> = (0..256).map(|i| (i as f64 * 0.1).sin() * 100.0).collect();
        let s = iaaft_surrogate(&x, 42, 100);

        // Sorted values should be identical (same marginal distribution)
        let mut x_sorted: Vec<f64> = x.clone();
        let mut s_sorted: Vec<f64> = s.clone();
        x_sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
        s_sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());

        for (a, b) in x_sorted.iter().zip(s_sorted.iter()) {
            assert!(
                (a - b).abs() < 1e-10,
                "Distribution not preserved: {a} vs {b}"
            );
        }
    }

    #[test]
    fn test_iaaft_different_seeds() {
        let x: Vec<f64> = (0..128).map(|i| (i as f64 * 0.1).sin()).collect();
        let s1 = iaaft_surrogate(&x, 1, 100);
        let s2 = iaaft_surrogate(&x, 2, 100);
        assert_ne!(
            s1, s2,
            "Different seeds should produce different surrogates"
        );
    }

    #[test]
    fn test_iaaft_small_input() {
        let x = vec![1.0, 2.0, 3.0];
        let s = iaaft_surrogate(&x, 42, 100);
        assert_eq!(s.len(), 3);
    }

    #[test]
    fn test_iaaft_preserves_amplitude_spectrum() {
        let n = 256;
        let x: Vec<f64> = (0..n)
            .map(|i| (i as f64 * 0.1).sin() + (i as f64 * 0.05).cos() * 0.5)
            .collect();
        let s = iaaft_surrogate(&x, 42, 200);

        // Compare amplitude spectra
        let mut planner = FftPlanner::new();
        let fft = planner.plan_fft_forward(n);

        let mut x_fft: Vec<Complex<f64>> = x.iter().map(|&v| Complex::new(v, 0.0)).collect();
        let mut s_fft: Vec<Complex<f64>> = s.iter().map(|&v| Complex::new(v, 0.0)).collect();
        fft.process(&mut x_fft);
        fft.process(&mut s_fft);

        let x_amps: Vec<f64> = x_fft.iter().map(|c| c.norm()).collect();
        let s_amps: Vec<f64> = s_fft.iter().map(|c| c.norm()).collect();

        // Amplitude spectra should be very close (not exact due to rank-ordering)
        let rmse: f64 = x_amps
            .iter()
            .zip(s_amps.iter())
            .map(|(a, b)| (a - b) * (a - b))
            .sum::<f64>()
            / n as f64;
        let rmse = rmse.sqrt();
        let max_amp = x_amps.iter().cloned().fold(0.0_f64, f64::max);

        // Relative RMSE should be small (< 5% of max amplitude)
        assert!(
            rmse / max_amp < 0.05,
            "Amplitude spectrum diverged: RMSE/max = {:.4}",
            rmse / max_amp
        );
    }
}
