scirs2-fft 0.4.2

//! Core scattering transform implementation
//!
//! Implements Mallat's scattering transform (2012):
//! - Zeroth order: S0 = x * phi (low-pass average)
//! - First order: S1 = |x * psi_{lambda1}| * phi
//! - Second order: S2 = ||x * psi_{lambda1}| * psi_{lambda2}| * phi
//!
//! The transform is translation invariant up to scale 2^J and Lipschitz
//! continuous to deformations, making it suitable as a feature extractor
//! for classification tasks.

use scirs2_core::numeric::Complex64;

use crate::error::{FFTError, FFTResult};
use crate::fft::{fft, ifft};

use super::filter_bank::{FilterBank, FilterBankConfig};

/// Configuration for the scattering transform.
#[derive(Debug, Clone)]
pub struct ScatteringConfig {
    /// Number of octaves
    pub j_max: usize,
    /// Quality factors per order (e.g., `[8, 1]`)
    pub quality_factors: Vec<usize>,
    /// Maximum scattering order (0, 1, or 2)
    pub max_order: usize,
    /// Whether to average (convolve with phi) the output
    pub average: bool,
    /// Subsampling factor for output (power of 2)
    pub oversampling: usize,
}

impl ScatteringConfig {
    /// Create a default scattering config.
    ///
    /// # Arguments
    /// * `j_max` - Number of octaves
    /// * `quality_factors` - Quality factors per order
    pub fn new(j_max: usize, quality_factors: Vec<usize>) -> Self {
        Self {
            j_max,
            quality_factors,
            max_order: 2,
            average: true,
            oversampling: 0,
        }
    }

    /// Set maximum scattering order.
    #[must_use]
    pub fn with_max_order(mut self, order: usize) -> Self {
        self.max_order = order.min(2);
        self
    }

    /// Set whether to average the output with the scaling function.
    #[must_use]
    pub fn with_average(mut self, average: bool) -> Self {
        self.average = average;
        self
    }

    /// Set oversampling factor.
    #[must_use]
    pub fn with_oversampling(mut self, oversampling: usize) -> Self {
        self.oversampling = oversampling;
        self
    }
}

/// Labels identifying which scattering order and path a coefficient belongs to.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum ScatteringOrder {
    /// Zeroth order: S0 = x * phi
    Zeroth,
    /// First order: S1 at wavelet index lambda1
    First {
        /// Index of the first-order wavelet filter.
        lambda1: usize,
    },
    /// Second order: S2 at wavelet indices (lambda1, lambda2)
    Second {
        /// Index of the first-order wavelet filter.
        lambda1: usize,
        /// Index of the second-order wavelet filter.
        lambda2: usize,
    },
}

/// A single scattering coefficient (time series or scalar).
#[derive(Debug, Clone)]
pub struct ScatteringCoefficients {
    /// Which order and path this coefficient belongs to
    pub order: ScatteringOrder,
    /// The coefficient values (time samples after subsampling)
    pub values: Vec<f64>,
}

/// Complete result of a scattering transform.
#[derive(Debug, Clone)]
pub struct ScatteringResult {
    /// All scattering coefficients, ordered: S0, then S1, then S2
    pub coefficients: Vec<ScatteringCoefficients>,
    /// Number of zeroth-order coefficients (always 1)
    pub num_zeroth: usize,
    /// Number of first-order paths
    pub num_first: usize,
    /// Number of second-order paths
    pub num_second: usize,
    /// Subsampled output length
    pub output_length: usize,
}

impl ScatteringResult {
    /// Get zeroth-order coefficients.
    pub fn zeroth_order(&self) -> &[ScatteringCoefficients] {
        &self.coefficients[..self.num_zeroth]
    }

    /// Get first-order coefficients.
    pub fn first_order(&self) -> &[ScatteringCoefficients] {
        &self.coefficients[self.num_zeroth..self.num_zeroth + self.num_first]
    }

    /// Get second-order coefficients.
    pub fn second_order(&self) -> &[ScatteringCoefficients] {
        &self.coefficients[self.num_zeroth + self.num_first..]
    }

    /// Flatten all coefficients into a single feature vector.
    pub fn flatten(&self) -> Vec<f64> {
        let mut result = Vec::new();
        for coeff in &self.coefficients {
            result.extend_from_slice(&coeff.values);
        }
        result
    }

    /// Total energy across all scattering coefficients.
    pub fn total_energy(&self) -> f64 {
        self.coefficients
            .iter()
            .flat_map(|c| c.values.iter())
            .map(|v| v * v)
            .sum()
    }
}

/// The scattering transform engine.
#[derive(Debug, Clone)]
pub struct ScatteringTransform {
    /// Scattering configuration
    config: ScatteringConfig,
    /// Pre-built filter bank
    filter_bank: FilterBank,
}

impl ScatteringTransform {
    /// Create a new scattering transform for signals of a given length.
    ///
    /// # Arguments
    /// * `config` - Scattering configuration
    /// * `signal_length` - Length of input signals
    pub fn new(config: ScatteringConfig, signal_length: usize) -> FFTResult<Self> {
        if signal_length == 0 {
            return Err(FFTError::ValueError(
                "signal_length must be positive".to_string(),
            ));
        }

        let fb_config =
            FilterBankConfig::new(config.j_max, config.quality_factors.clone(), signal_length);
        let filter_bank = FilterBank::new(fb_config)?;

        Ok(Self {
            config,
            filter_bank,
        })
    }

    /// Access the underlying filter bank.
    pub fn filter_bank(&self) -> &FilterBank {
        &self.filter_bank
    }

    /// Compute the scattering transform of a real-valued signal.
    ///
    /// Returns scattering coefficients organized by order.
    pub fn transform(&self, signal: &[f64]) -> FFTResult<ScatteringResult> {
        if signal.is_empty() {
            return Err(FFTError::ValueError(
                "Input signal must not be empty".to_string(),
            ));
        }

        let fft_size = self.filter_bank.fft_size;

        // Pad signal to FFT size
        let mut padded = vec![0.0_f64; fft_size];
        let copy_len = signal.len().min(fft_size);
        padded[..copy_len].copy_from_slice(&signal[..copy_len]);

        // Compute FFT of input
        let x_hat = fft(&padded, Some(fft_size))?;

        // Subsampling factor
        let subsample = if self.config.average {
            let base = 2_usize.pow(self.config.j_max as u32);
            base >> self.config.oversampling.min(self.config.j_max)
        } else {
            1
        };
        let output_length = fft_size.div_ceil(subsample);

        let mut coefficients = Vec::new();

        let mut num_first = 0;
        let mut num_second = 0;

        // --- Zeroth order: S0 = x * phi ---
        let s0 = convolve_and_subsample(&x_hat, &self.filter_bank.phi, fft_size, subsample)?;
        coefficients.push(ScatteringCoefficients {
            order: ScatteringOrder::Zeroth,
            values: s0,
        });
        let num_zeroth = 1;

        if self.config.max_order == 0 {
            return Ok(ScatteringResult {
                coefficients,
                num_zeroth,
                num_first,
                num_second,
                output_length,
            });
        }

        // --- First order: S1 = |x * psi_{lambda1}| * phi ---
        let first_order_wavelets = self
            .filter_bank
            .wavelets
            .first()
            .ok_or_else(|| FFTError::ComputationError("No first-order wavelets".to_string()))?;

        // Store U1 (unaveraged first-order) for second-order computation
        let mut u1_hats: Vec<Vec<Complex64>> = Vec::new();

        for (lambda1, wavelet) in first_order_wavelets.iter().enumerate() {
            // x * psi_{lambda1} in frequency domain
            let convolved: Vec<Complex64> = x_hat
                .iter()
                .zip(wavelet.freq_response.iter())
                .map(|(x, w)| x * w)
                .collect();

            // IFFT to get time-domain convolution
            let u1_time = ifft(&convolved, None)?;

            // Modulus: |x * psi_{lambda1}|
            let u1_mod: Vec<f64> = u1_time.iter().map(|c| c.norm()).collect();

            // Store FFT of modulus for second-order computation
            if self.config.max_order >= 2 {
                let u1_mod_hat = fft(&u1_mod, Some(fft_size))?;
                u1_hats.push(u1_mod_hat);
            }

            // Average with phi: |x * psi_{lambda1}| * phi
            let u1_mod_hat_for_avg = if self.config.max_order >= 2 {
                // Already computed above; reuse
                u1_hats.last().ok_or_else(|| {
                    FFTError::ComputationError("u1_hats should not be empty".to_string())
                })?
            } else {
                // Compute just for averaging
                &fft(&u1_mod, Some(fft_size))?
            };

            let s1 = convolve_and_subsample(
                u1_mod_hat_for_avg,
                &self.filter_bank.phi,
                fft_size,
                subsample,
            )?;

            coefficients.push(ScatteringCoefficients {
                order: ScatteringOrder::First { lambda1 },
                values: s1,
            });
            num_first += 1;
        }

        if self.config.max_order < 2 {
            return Ok(ScatteringResult {
                coefficients,
                num_zeroth,
                num_first,
                num_second,
                output_length,
            });
        }

        // --- Second order: S2 = ||x * psi_{lambda1}| * psi_{lambda2}| * phi ---
        // Only for lambda2 > lambda1 (coarser scale than lambda1)
        let second_order_wavelets = if self.filter_bank.wavelets.len() > 1 {
            &self.filter_bank.wavelets[1]
        } else {
            // Use first-order wavelets if no separate second-order bank
            &self.filter_bank.wavelets[0]
        };

        for (lambda1, u1_hat) in u1_hats.iter().enumerate() {
            for (lambda2, wavelet2) in second_order_wavelets.iter().enumerate() {
                // Only compute when lambda2 represents a coarser scale
                // For second-order wavelets with different Q, compare scale indices
                let first_scale = if !first_order_wavelets.is_empty() {
                    first_order_wavelets[lambda1].j
                } else {
                    0
                };
                let second_scale = wavelet2.j;

                if second_scale <= first_scale {
                    continue;
                }

                // |U1| * psi_{lambda2}
                let convolved2: Vec<Complex64> = u1_hat
                    .iter()
                    .zip(wavelet2.freq_response.iter())
                    .map(|(u, w)| u * w)
                    .collect();

                let u2_time = ifft(&convolved2, None)?;

                // Modulus
                let u2_mod: Vec<f64> = u2_time.iter().map(|c| c.norm()).collect();

                // Average with phi
                let u2_mod_hat = fft(&u2_mod, Some(fft_size))?;
                let s2 = convolve_and_subsample(
                    &u2_mod_hat,
                    &self.filter_bank.phi,
                    fft_size,
                    subsample,
                )?;

                coefficients.push(ScatteringCoefficients {
                    order: ScatteringOrder::Second { lambda1, lambda2 },
                    values: s2,
                });
                num_second += 1;
            }
        }

        Ok(ScatteringResult {
            coefficients,
            num_zeroth,
            num_first,
            num_second,
            output_length,
        })
    }

    /// Compute the scattering transform and return only the feature vector.
    pub fn features(&self, signal: &[f64]) -> FFTResult<Vec<f64>> {
        let result = self.transform(signal)?;
        Ok(result.flatten())
    }
}

/// Multiply two spectra element-wise, IFFT, take real part, and subsample.
fn convolve_and_subsample(
    x_hat: &[Complex64],
    filter_hat: &[Complex64],
    fft_size: usize,
    subsample: usize,
) -> FFTResult<Vec<f64>> {
    // Pointwise multiplication in frequency domain
    let product: Vec<Complex64> = x_hat
        .iter()
        .zip(filter_hat.iter())
        .map(|(x, f)| x * f)
        .collect();

    // IFFT
    let time_domain = ifft(&product, None)?;

    // Subsample and take real part
    let output_len = fft_size.div_ceil(subsample);
    let mut result = Vec::with_capacity(output_len);
    for i in 0..output_len {
        let idx = i * subsample;
        if idx < time_domain.len() {
            result.push(time_domain[idx].re);
        } else {
            result.push(0.0);
        }
    }

    Ok(result)
}

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

    #[test]
    fn test_scattering_basic() {
        let config = ScatteringConfig::new(3, vec![2, 1]);
        let st = ScatteringTransform::new(config, 256)
            .expect("scattering transform creation should succeed");

        // Simple sine wave
        let signal: Vec<f64> = (0..256)
            .map(|i| (2.0 * PI * 10.0 * i as f64 / 256.0).sin())
            .collect();

        let result = st.transform(&signal).expect("transform should succeed");

        assert_eq!(result.num_zeroth, 1);
        assert!(result.num_first > 0);
        // Second-order should exist but may be 0 if no valid lambda2 > lambda1
    }

    #[test]
    fn test_translation_invariance() {
        // Translation invariance is demonstrated by comparing the total energy
        // of first-order scattering coefficients for a signal and its circular shift.
        // The scattering transform averages over a window of 2^J samples, so
        // circular translations should produce similar total first-order energies.
        let config = ScatteringConfig::new(3, vec![4, 1]).with_max_order(1);
        let n = 512;
        let st = ScatteringTransform::new(config, n)
            .expect("scattering transform creation should succeed");

        // Signal: a Gaussian pulse
        let mut signal1 = vec![0.0; n];
        for i in 0..n {
            let t = (i as f64 - 128.0) / 20.0;
            signal1[i] = (-0.5 * t * t).exp();
        }

        // Circularly shifted version
        let shift = 64;
        let mut signal2 = vec![0.0; n];
        for i in 0..n {
            let src = (i + n - shift) % n;
            signal2[i] = signal1[src];
        }

        let r1 = st.transform(&signal1).expect("transform should succeed");
        let r2 = st.transform(&signal2).expect("transform should succeed");

        // Compare total energy of first-order coefficients
        // Each S1 path is |x * psi| * phi, so translating x circularly
        // should give similar energies per path
        let s1_energies_1: Vec<f64> = r1
            .first_order()
            .iter()
            .map(|c| c.values.iter().map(|v| v * v).sum::<f64>())
            .collect();
        let s1_energies_2: Vec<f64> = r2
            .first_order()
            .iter()
            .map(|c| c.values.iter().map(|v| v * v).sum::<f64>())
            .collect();

        let total_e1: f64 = s1_energies_1.iter().sum();
        let total_e2: f64 = s1_energies_2.iter().sum();

        if total_e1 > 1e-15 {
            let rel_error = ((total_e1 - total_e2) / total_e1).abs();
            assert!(
                rel_error < 0.3,
                "First-order total energy should be approximately translation invariant, \
                 rel_error={:.4} (e1={:.4}, e2={:.4})",
                rel_error,
                total_e1,
                total_e2
            );
        }
    }

    #[test]
    fn test_output_dimensions() {
        let j = 3;
        let q1 = 4;
        let q2 = 1;
        let config = ScatteringConfig::new(j, vec![q1, q2]);
        let n = 256;
        let st = ScatteringTransform::new(config, n)
            .expect("scattering transform creation should succeed");

        let signal: Vec<f64> = (0..n)
            .map(|i| (2.0 * PI * 5.0 * i as f64 / n as f64).sin())
            .collect();

        let result = st.transform(&signal).expect("transform should succeed");

        // First order: J * Q1 = 3 * 4 = 12 paths
        assert_eq!(result.num_first, j * q1);

        // Second order depends on the lambda2 > lambda1 condition
        // With Q2=1, we have J*Q2=3 second-order wavelets
        // For each first-order path, only coarser second-order wavelets apply
        // num_second can be any non-negative value depending on scale ordering
        let _ = result.num_second;

        // All coefficients should have the same output length
        let expected_len = result.output_length;
        for coeff in &result.coefficients {
            assert_eq!(
                coeff.values.len(),
                expected_len,
                "coefficient output length mismatch"
            );
        }
    }

    #[test]
    fn test_energy_approximate_preservation() {
        let config = ScatteringConfig::new(3, vec![4, 1]);
        let n = 256;
        let st = ScatteringTransform::new(config, n)
            .expect("scattering transform creation should succeed");

        let signal: Vec<f64> = (0..n)
            .map(|i| {
                let t = i as f64 / n as f64;
                (2.0 * PI * 8.0 * t).sin() + 0.5 * (2.0 * PI * 32.0 * t).cos()
            })
            .collect();

        let input_energy: f64 = signal.iter().map(|v| v * v).sum();
        let result = st.transform(&signal).expect("transform should succeed");
        let scatter_energy = result.total_energy();

        // Scattering energy should be bounded by input energy
        // Due to subsampling and the scattering inequality, scatter_energy <= input_energy
        // but a significant fraction should be preserved
        assert!(scatter_energy > 0.0, "scattering energy should be positive");
    }

    #[test]
    fn test_sine_wave_first_order() {
        let config = ScatteringConfig::new(4, vec![8]).with_max_order(1);
        let n = 1024;
        let st = ScatteringTransform::new(config, n)
            .expect("scattering transform creation should succeed");

        // Pure sine wave at a known frequency
        let freq = 20.0; // cycles per signal length
        let signal: Vec<f64> = (0..n)
            .map(|i| (2.0 * PI * freq * i as f64 / n as f64).sin())
            .collect();

        let result = st.transform(&signal).expect("transform should succeed");

        // First-order coefficients: the strongest response should come from
        // the wavelet whose center frequency is closest to the sine frequency.
        let first = result.first_order();
        assert!(!first.is_empty(), "should have first-order coefficients");

        // Find the path with maximum energy
        let max_path = first
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| {
                let ea: f64 = a.values.iter().map(|v| v * v).sum();
                let eb: f64 = b.values.iter().map(|v| v * v).sum();
                ea.partial_cmp(&eb).unwrap_or(std::cmp::Ordering::Equal)
            })
            .map(|(idx, _)| idx);

        assert!(max_path.is_some(), "should find a path with maximum energy");
    }

    #[test]
    fn test_zeroth_order_only() {
        let config = ScatteringConfig::new(3, vec![4]).with_max_order(0);
        let n = 128;
        let st = ScatteringTransform::new(config, n)
            .expect("scattering transform creation should succeed");

        let signal: Vec<f64> = (0..n).map(|i| i as f64 / n as f64).collect();
        let result = st.transform(&signal).expect("transform should succeed");

        assert_eq!(result.num_zeroth, 1);
        assert_eq!(result.num_first, 0);
        assert_eq!(result.num_second, 0);
    }

    #[test]
    fn test_empty_signal_error() {
        let config = ScatteringConfig::new(3, vec![4]);
        let st = ScatteringTransform::new(config, 128)
            .expect("scattering transform creation should succeed");

        let result = st.transform(&[]);
        assert!(result.is_err());
    }
}