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

/// Compute the 3D Riesz transform of a scalar field.
///
/// Input: real-valued scalar field on an N×N×N grid (e.g., SDF values).
/// Output: three vector-valued fields (R₁f, R₂f, R₃f), each N×N×N.
///
/// In the frequency domain:
///   R̂ⱼ(ξ) = -i · (ξⱼ / |ξ|) · f̂(ξ)
///
/// This is a phase rotation proportional to the direction cosine
/// of each frequency component. Exactly what CORDIC rotation does.
pub struct RieszTransform {
    pub r1: Vec<f64>,
    pub r2: Vec<f64>,
    pub r3: Vec<f64>,
    pub n: usize,
}

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

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

        fft3d(&mut spectrum, n);

        let mut r1_spec = vec![Complex::new(0.0, 0.0); n * n * n];
        let mut r2_spec = vec![Complex::new(0.0, 0.0); n * n * n];
        let mut r3_spec = vec![Complex::new(0.0, 0.0); n * n * n];

        for iz in 0..n {
            for iy in 0..n {
                for ix in 0..n {
                    let idx = iz * n * n + iy * n + ix;

                    // Frequency coordinates (centered)
                    let fx = if ix <= n / 2 { ix as f64 } else { ix as f64 - n as f64 };
                    let fy = if iy <= n / 2 { iy as f64 } else { iy as f64 - n as f64 };
                    let fz = if iz <= n / 2 { iz as f64 } else { iz as f64 - n as f64 };

                    let mag = (fx * fx + fy * fy + fz * fz).sqrt();

                    if mag < 1e-10 {
                        // DC component — Riesz transform of DC is zero
                        continue;
                    }

                    // R̂ⱼ(ξ) = -i · (ξⱼ / |ξ|) · f̂(ξ)
                    // Multiplying by -i rotates: (a + bi) * (-i) = b - ai
                    let f_hat = spectrum[idx];
                    let neg_i_f = Complex::new(f_hat.im, -f_hat.re);

                    r1_spec[idx] = neg_i_f * (fx / mag);
                    r2_spec[idx] = neg_i_f * (fy / mag);
                    r3_spec[idx] = neg_i_f * (fz / mag);
                }
            }
        }

        ifft3d(&mut r1_spec, n);
        ifft3d(&mut r2_spec, n);
        ifft3d(&mut r3_spec, n);

        RieszTransform {
            r1: r1_spec.iter().map(|c| c.re).collect(),
            r2: r2_spec.iter().map(|c| c.re).collect(),
            r3: r3_spec.iter().map(|c| c.re).collect(),
            n,
        }
    }
}
