use crate::fft3d::{fft3d, ifft3d};
use rustfft::num_complex::Complex;

/// 3D spatial cepstrum: detects periodic structures in a scalar field.
///
/// Cepstrum = IFFT(log(|FFT(f)|))
///
/// Peaks in the cepstrum correspond to periodicities in the original field:
/// screw threads, gear teeth, bolt patterns, array features, lattice structures.
///
/// The position of a peak gives the period vector (direction + spacing).
/// The height of a peak gives the strength of the periodicity.
pub struct Cepstrum {
    pub data: Vec<f64>,
    pub n: usize,
}

#[derive(Debug, Clone, Copy)]
pub struct PeriodicFeature {
    /// Grid offset of the period (direction and spacing).
    pub dx: i32,
    pub dy: i32,
    pub dz: i32,
    /// Cepstral magnitude (strength of periodicity).
    pub strength: f64,
}

impl Cepstrum {
    pub fn compute(field: &[f64], n: usize) -> Self {
        assert_eq!(field.len(), n * n * n);

        let mut spectrum: Vec<Complex<f64>> = field.iter()
            .map(|&v| Complex::new(v, 0.0))
            .collect();

        fft3d(&mut spectrum, n);

        // log(|FFT|) — use log of magnitude
        let floor = 1e-10;
        for val in spectrum.iter_mut() {
            let mag = val.norm().max(floor);
            *val = Complex::new(mag.ln(), 0.0);
        }

        ifft3d(&mut spectrum, n);

        Cepstrum {
            data: spectrum.iter().map(|c| c.re).collect(),
            n,
        }
    }

    /// Find the strongest periodic features.
    /// Excludes the origin (DC component) and its immediate neighbors.
    pub fn detect_periodicities(&self, count: usize) -> Vec<PeriodicFeature> {
        let n = self.n;
        let half = n / 2;
        let mut peaks: Vec<PeriodicFeature> = Vec::new();

        for iz in 0..n {
            for iy in 0..n {
                for ix in 0..n {
                    // Skip origin region (not a periodicity)
                    let dx = if ix <= half { ix as i32 } else { ix as i32 - n as i32 };
                    let dy = if iy <= half { iy as i32 } else { iy as i32 - n as i32 };
                    let dz = if iz <= half { iz as i32 } else { iz as i32 - n as i32 };

                    let dist_sq = dx * dx + dy * dy + dz * dz;
                    if dist_sq <= 4 { continue; } // skip DC neighborhood

                    let val = self.data[iz * n * n + iy * n + ix].abs();
                    peaks.push(PeriodicFeature {
                        dx, dy, dz,
                        strength: val,
                    });
                }
            }
        }

        peaks.sort_by(|a, b| b.strength.partial_cmp(&a.strength).unwrap_or(std::cmp::Ordering::Equal));
        peaks.truncate(count);

        // Deduplicate: (dx,dy,dz) and (-dx,-dy,-dz) are the same periodicity
        let mut deduped: Vec<PeriodicFeature> = Vec::new();
        for p in &peaks {
            let is_dup = deduped.iter().any(|q| {
                (q.dx == -p.dx && q.dy == -p.dy && q.dz == -p.dz)
                || (q.dx == p.dx && q.dy == p.dy && q.dz == p.dz)
            });
            if !is_dup {
                deduped.push(*p);
            }
        }

        deduped
    }
}
