oxifft 0.4.0

//! Non-uniform FFT (NUFFT) implementation.
//!
//! This module provides FFT for non-equispaced (non-uniform) data points,
//! which is essential for applications like:
//! - MRI reconstruction
//! - Radio astronomy
//! - Seismic imaging
//! - Spectral analysis of irregularly sampled signals
//!
//! # NUFFT Types
//!
//! - **Type 1 (Adjoint)**: Non-uniform points → uniform grid
//! - **Type 2 (Forward)**: Uniform grid → non-uniform points
//! - **Type 3**: Non-uniform → non-uniform
//!
//! # Algorithm
//!
//! Uses the Gaussian gridding approach:
//! 1. Spread non-uniform data to oversampled grid (convolution with kernel)
//! 2. Apply standard FFT
//! 3. Deconvolve to correct for spreading kernel
//!
//! # Example
//!
//! ```
//! use oxifft::nufft::{Nufft, NufftType};
//! use oxifft::Complex;
//!
//! // Non-uniform sample locations in [-π, π]
//! let x = vec![-2.0, -0.5, 0.3, 1.5, 2.8];
//! let values = vec![Complex::new(1.0_f64, 0.0); 5];
//!
//! // Create NUFFT plan
//! let plan = Nufft::new(NufftType::Type1, 64, &x, 1e-6).expect("valid NUFFT plan");
//!
//! // Execute: non-uniform → uniform grid
//! let result = plan.execute(&values).expect("NUFFT execution");
//! # let _ = result;
//! ```

use crate::api::{Direction, Flags, Plan};
use crate::kernel::{Complex, Float};
use crate::prelude::*;

pub mod nufft2d;
pub mod nufft3d;

pub use nufft2d::{nufft2d_type1, nufft2d_type2};
pub use nufft3d::nufft3d_type1;

/// NUFFT type specifying the direction of transformation.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum NufftType {
    /// Type 1: Non-uniform to uniform (adjoint NUFFT).
    /// Given values at non-uniform points, compute uniform Fourier coefficients.
    Type1,
    /// Type 2: Uniform to non-uniform (forward NUFFT).
    /// Given uniform Fourier coefficients, evaluate at non-uniform points.
    Type2,
    /// Type 3: Non-uniform to non-uniform.
    /// Transform between two sets of non-uniform points.
    Type3,
}

/// NUFFT configuration options.
#[derive(Debug, Clone, Copy)]
pub struct NufftOptions {
    /// Oversampling factor (typically 2.0).
    pub oversampling: f64,
    /// Kernel width in grid points (typically 4-12).
    pub kernel_width: usize,
    /// Target relative tolerance.
    pub tolerance: f64,
    /// Use multi-threaded spreading.
    pub threaded: bool,
}

impl Default for NufftOptions {
    fn default() -> Self {
        Self {
            oversampling: 2.0,
            kernel_width: 6,
            tolerance: 1e-6,
            threaded: true,
        }
    }
}

/// NUFFT error types.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum NufftError {
    /// Invalid input size.
    InvalidSize(usize),
    /// Points out of range [-π, π].
    PointsOutOfRange,
    /// FFT planning failed.
    PlanFailed,
    /// Execution failed.
    ExecutionFailed(String),
    /// Invalid tolerance.
    InvalidTolerance,
}

impl core::fmt::Display for NufftError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::InvalidSize(n) => write!(f, "Invalid NUFFT size: {n}"),
            Self::PointsOutOfRange => write!(f, "Non-uniform points must be in [-π, π]"),
            Self::PlanFailed => write!(f, "Failed to create FFT plan"),
            Self::ExecutionFailed(msg) => write!(f, "NUFFT execution failed: {msg}"),
            Self::InvalidTolerance => write!(f, "Tolerance must be positive"),
        }
    }
}

/// Result type for NUFFT operations.
pub type NufftResult<T> = Result<T, NufftError>;

/// Non-uniform FFT plan.
///
/// Precomputes spreading/interpolation coefficients for efficient repeated use.
#[allow(clippy::struct_field_names)] // reason: fields named by mathematical role (nufft_type, n_uniform, etc.); renaming would obscure the NUFFT algorithm structure
pub struct Nufft<T: Float> {
    /// NUFFT type.
    nufft_type: NufftType,
    /// Number of uniform grid points.
    n_uniform: usize,
    /// Number of non-uniform points.
    n_nonuniform: usize,
    /// Oversampled grid size.
    n_oversampled: usize,
    /// Non-uniform point locations (normalized to [0, 2π]).
    points: Vec<f64>,
    /// Precomputed spreading coefficients for each non-uniform point.
    spread_coeffs: Vec<Vec<(usize, T)>>,
    /// Deconvolution factors.
    deconv_factors: Vec<Complex<T>>,
    /// Internal forward FFT plan (oversampled grid, used by Type 1 spreading).
    fft_plan: Option<Plan<T>>,
    /// Internal inverse FFT plan (oversampled grid, used by Type 2 interpolation).
    ///
    /// Precomputed once at construction time (rather than per-`type2()` call)
    /// so repeated Type 2 execution does not pay planning cost on every call.
    inv_fft_plan: Option<Plan<T>>,
    /// Planning flags used to build `fft_plan` / `inv_fft_plan`.
    flags: Flags,
    /// Options.
    options: NufftOptions,
}

impl<T: Float> Nufft<T> {
    /// Create a new NUFFT plan.
    ///
    /// # Arguments
    ///
    /// * `nufft_type` - Type of NUFFT (1, 2, or 3)
    /// * `n_uniform` - Number of uniform grid points
    /// * `points` - Non-uniform point locations in [-π, π]
    /// * `tolerance` - Target relative accuracy
    ///
    /// # Returns
    ///
    /// NUFFT plan or error.
    ///
    /// # Errors
    ///
    /// Returns error if size is zero, tolerance is non-positive, or points are out of range.
    ///
    /// # Examples
    ///
    /// ```
    /// use oxifft::{Complex, Nufft, NufftType};
    ///
    /// // Type 1 NUFFT: 4 non-uniform points → 8 uniform Fourier modes
    /// let points = vec![-1.0_f64, -0.3, 0.2, 0.8];
    /// let plan = Nufft::<f64>::new(NufftType::Type1, 8, &points, 1e-6)
    ///     .expect("NUFFT plan creation failed");
    /// let values: Vec<Complex<f64>> = points.iter()
    ///     .map(|&x| Complex::new(x.cos(), 0.0))
    ///     .collect();
    /// let modes = plan.type1(&values).expect("type1 execution failed");
    /// // Output has n_uniform = 8 modes
    /// assert_eq!(modes.len(), 8);
    /// ```
    ///
    /// Uses [`Flags::ESTIMATE`] for the internal FFT plans (fast planning, no
    /// benchmarking). Use [`Nufft::with_flags`] or
    /// [`Nufft::with_options_and_flags`] to request a different planning
    /// strategy (e.g. `Flags::MEASURE` for repeated executions of the same
    /// plan where planning cost is amortized).
    pub fn new(
        nufft_type: NufftType,
        n_uniform: usize,
        points: &[f64],
        tolerance: f64,
    ) -> NufftResult<Self> {
        let options = NufftOptions {
            tolerance,
            ..Default::default()
        };
        Self::with_options(nufft_type, n_uniform, points, &options)
    }

    /// Create NUFFT plan with custom options.
    ///
    /// Uses [`Flags::ESTIMATE`] for the internal FFT plans. Use
    /// [`Nufft::with_options_and_flags`] to request a different planning
    /// strategy.
    ///
    /// # Errors
    ///
    /// Returns error if size is zero, tolerance is non-positive, or points are out of range.
    pub fn with_options(
        nufft_type: NufftType,
        n_uniform: usize,
        points: &[f64],
        options: &NufftOptions,
    ) -> NufftResult<Self> {
        Self::with_options_and_flags(nufft_type, n_uniform, points, options, Flags::ESTIMATE)
    }

    /// Create a new NUFFT plan with an explicit FFT planning strategy.
    ///
    /// Equivalent to [`Nufft::new`] but lets the caller choose the [`Flags`]
    /// used for the internal oversampled-grid FFT plans (e.g. `Flags::MEASURE`
    /// to spend more time planning in exchange for faster repeated execution).
    ///
    /// # Errors
    ///
    /// Returns error if size is zero, tolerance is non-positive, or points are out of range.
    pub fn with_flags(
        nufft_type: NufftType,
        n_uniform: usize,
        points: &[f64],
        tolerance: f64,
        flags: Flags,
    ) -> NufftResult<Self> {
        let options = NufftOptions {
            tolerance,
            ..Default::default()
        };
        Self::with_options_and_flags(nufft_type, n_uniform, points, &options, flags)
    }

    /// Create NUFFT plan with custom options and an explicit FFT planning strategy.
    ///
    /// Both the forward (Type 1 spreading) and inverse (Type 2 interpolation)
    /// oversampled-grid FFT plans are built once here, using the same `flags`,
    /// so that repeated `type1`/`type2` execution never re-plans internally.
    ///
    /// # Errors
    ///
    /// Returns error if size is zero, tolerance is non-positive, or points are out of range.
    pub fn with_options_and_flags(
        nufft_type: NufftType,
        n_uniform: usize,
        points: &[f64],
        options: &NufftOptions,
        flags: Flags,
    ) -> NufftResult<Self> {
        if n_uniform == 0 {
            return Err(NufftError::InvalidSize(0));
        }
        if options.tolerance <= 0.0 {
            return Err(NufftError::InvalidTolerance);
        }

        // Compute kernel width from tolerance and oversampling ratio.
        // The oversampling ratio affects how large the kernel needs to be to
        // achieve the desired accuracy (smaller oversampling needs wider kernel).
        let kernel_width = compute_kernel_width(
            options.tolerance,
            options.oversampling,
            options.kernel_width,
        );

        // Compute oversampled grid size
        let n_oversampled = ((n_uniform as f64) * options.oversampling).ceil() as usize;
        let n_oversampled = next_smooth_number(n_oversampled);

        // Normalize points to [0, 2π] and validate
        let mut normalized_points = Vec::with_capacity(points.len());
        for &p in points {
            if !(-core::f64::consts::PI..=core::f64::consts::PI).contains(&p) {
                return Err(NufftError::PointsOutOfRange);
            }
            // Shift from [-π, π] to [0, 2π]
            normalized_points.push(p + core::f64::consts::PI);
        }

        // Precompute spreading coefficients
        let spread_coeffs =
            precompute_spreading_coeffs(&normalized_points, n_oversampled, kernel_width);

        // Precompute deconvolution factors
        let deconv_factors = precompute_deconv_factors(n_uniform, n_oversampled, kernel_width);

        // Create forward and inverse FFT plans once, using the same flags, so
        // that Type 1 (forward) and Type 2 (inverse) execution never re-plan
        // on every call.
        let fft_plan = Plan::dft_1d(n_oversampled, Direction::Forward, flags);
        let inv_fft_plan = Plan::dft_1d(n_oversampled, Direction::Backward, flags);

        Ok(Self {
            nufft_type,
            n_uniform,
            n_nonuniform: points.len(),
            n_oversampled,
            points: normalized_points,
            spread_coeffs,
            deconv_factors,
            fft_plan,
            inv_fft_plan,
            flags,
            options: NufftOptions {
                kernel_width,
                ..*options
            },
        })
    }

    /// Execute Type 1 NUFFT: non-uniform → uniform.
    ///
    /// Given values at non-uniform points, compute uniform Fourier coefficients.
    ///
    /// # Errors
    ///
    /// Returns error if input length doesn't match the number of non-uniform points.
    pub fn type1(&self, values: &[Complex<T>]) -> NufftResult<Vec<Complex<T>>> {
        if values.len() != self.n_nonuniform {
            return Err(NufftError::ExecutionFailed(format!(
                "Expected {} values, got {}",
                self.n_nonuniform,
                values.len()
            )));
        }

        // Step 1: Spread non-uniform data to oversampled grid
        let mut grid = vec![Complex::<T>::zero(); self.n_oversampled];
        self.spread_to_grid(values, &mut grid);

        // Step 2: FFT on oversampled grid
        let mut fft_result = vec![Complex::<T>::zero(); self.n_oversampled];
        if let Some(ref plan) = self.fft_plan {
            plan.execute(&grid, &mut fft_result);
        } else {
            return Err(NufftError::PlanFailed);
        }

        // Step 3: Deconvolve and extract frequencies in math order.
        //
        // The output convention (matching the dense NDFT / FINUFFT Type 1) is:
        //   result[k] corresponds to frequency  freq = k − N/2
        //   k=0      → freq = −N/2   (most-negative)
        //   k=N/2    → freq = 0      (DC)
        //   k=N−1    → freq = N/2−1  (most-positive)
        //
        // In the oversampled FFT result, frequency `freq` lives at:
        //   grid_idx = freq               if freq ≥ 0
        //   grid_idx = n_oversampled+freq if freq < 0
        //
        // The deconv_factors array is in FFT order (index `d` → freq `d` if
        // d < N/2, freq `d−N` if d ≥ N/2), so the FFT-order index for `freq` is:
        //   deconv_idx = freq               if freq ≥ 0
        //   deconv_idx = n_uniform + freq   if freq < 0
        let mut result = Vec::with_capacity(self.n_uniform);

        for k in 0..self.n_uniform {
            let (grid_idx, deconv_idx) =
                centered_freq_indices(k, self.n_uniform, self.n_oversampled);
            result.push(fft_result[grid_idx] * self.deconv_factors[deconv_idx]);
        }

        Ok(result)
    }

    /// Execute Type 2 NUFFT: uniform → non-uniform.
    ///
    /// Given uniform Fourier coefficients, evaluate at non-uniform points.
    ///
    /// # Errors
    ///
    /// Returns error if coefficient length doesn't match the uniform grid size.
    pub fn type2(&self, coeffs: &[Complex<T>]) -> NufftResult<Vec<Complex<T>>> {
        if coeffs.len() != self.n_uniform {
            return Err(NufftError::ExecutionFailed(format!(
                "Expected {} coefficients, got {}",
                self.n_uniform,
                coeffs.len()
            )));
        }

        // Step 1: Deconvolve coefficients and scatter into oversampled grid.
        //
        // Input convention (matching the dense NDFT / FINUFFT Type 2):
        //   coeffs[k] is the Fourier coefficient for frequency  freq = k − N/2
        //   k=0   → freq = −N/2  (most-negative)
        //   k=N/2 → freq = 0     (DC)
        //
        // In the oversampled grid, frequency `freq` is placed at:
        //   grid_idx = freq               if freq ≥ 0
        //   grid_idx = n_oversampled+freq if freq < 0
        //
        // deconv_factors is in FFT order; the FFT-order index for `freq` is:
        //   deconv_idx = freq               if freq ≥ 0
        //   deconv_idx = n_uniform + freq   if freq < 0
        //
        // Type 2 deconvolution differs from Type 1 by a factor of n_oversampled.
        // After IFFT (unnormalized) + 1/n_os normalization + kernel interpolation,
        // the output picks up an extra factor of Ψ̂(freq)/n_os from the
        // interpolation step.  To cancel this, the grid coefficient must be
        // n_os / Ψ̂(freq), not 1 / Ψ̂(freq).  Since deconv_factors already stores
        // 1/Ψ̂, we multiply by n_os here.
        let mut grid = vec![Complex::<T>::zero(); self.n_oversampled];
        let n_os_scale = T::from_usize(self.n_oversampled);

        for (k, &coeff) in coeffs.iter().enumerate() {
            let (grid_idx, deconv_idx) =
                centered_freq_indices(k, self.n_uniform, self.n_oversampled);

            // Multiply deconv factor by n_os to account for the IFFT normalization
            // that is undone by the interpolation (kernel re-applies a factor of ~1/n_os
            // relative to what Type 1 spreading adds, so Type 2 needs n_os × the
            // Type 1 deconv factor).
            let scaled_deconv = Complex::new(
                self.deconv_factors[deconv_idx].re * n_os_scale,
                self.deconv_factors[deconv_idx].im * n_os_scale,
            );
            grid[grid_idx] = coeff * scaled_deconv;
        }

        // Step 2: Inverse FFT (uses the plan precomputed at construction time)
        let mut ifft_result = vec![Complex::<T>::zero(); self.n_oversampled];
        if let Some(ref inv_plan) = self.inv_fft_plan {
            inv_plan.execute(&grid, &mut ifft_result);
        } else {
            return Err(NufftError::PlanFailed);
        }

        // Normalize
        let scale = T::ONE / T::from_usize(self.n_oversampled);
        for c in &mut ifft_result {
            *c = Complex::new(c.re * scale, c.im * scale);
        }

        // Step 3: Interpolate at non-uniform points
        let result = self.interpolate_from_grid(&ifft_result);

        Ok(result)
    }

    /// Execute NUFFT based on the configured type.
    ///
    /// # Errors
    ///
    /// Returns error for Type3 (use `execute_type3` instead) or if input validation fails.
    pub fn execute(&self, input: &[Complex<T>]) -> NufftResult<Vec<Complex<T>>> {
        match self.nufft_type {
            NufftType::Type1 => self.type1(input),
            NufftType::Type2 => self.type2(input),
            NufftType::Type3 => {
                // Type 3 = Type 1 followed by Type 2 (with different target points)
                // For simplicity, we require separate source/target points
                Err(NufftError::ExecutionFailed(
                    "Type 3 requires separate execute_type3 call".into(),
                ))
            }
        }
    }

    /// Execute Type 3 NUFFT: non-uniform → non-uniform.
    ///
    /// # Arguments
    ///
    /// * `values` - Values at source points (set during plan creation)
    /// * `target_points` - Target non-uniform points in [-π, π]
    ///
    /// # Errors
    ///
    /// Returns error if input validation fails or target points are out of range.
    pub fn execute_type3(
        &self,
        values: &[Complex<T>],
        target_points: &[f64],
    ) -> NufftResult<Vec<Complex<T>>> {
        // Type 3 = Type 1 to uniform grid, then Type 2 to target points
        // First do Type 1
        let uniform_coeffs = self.type1(values)?;

        // Create Type 2 plan for target points, propagating this plan's
        // options and FFT planning flags so Type 3 behaves consistently with
        // the Type 1 half of the computation.
        let type2_plan = Self::with_options_and_flags(
            NufftType::Type2,
            self.n_uniform,
            target_points,
            &self.options,
            self.flags,
        )?;

        // Execute Type 2
        type2_plan.type2(&uniform_coeffs)
    }

    /// Spread non-uniform values to the oversampled grid.
    fn spread_to_grid(&self, values: &[Complex<T>], grid: &mut [Complex<T>]) {
        for (j, &val) in values.iter().enumerate() {
            for &(idx, weight) in &self.spread_coeffs[j] {
                grid[idx] = grid[idx] + Complex::new(val.re * weight, val.im * weight);
            }
        }
    }

    /// Interpolate from grid at non-uniform points.
    fn interpolate_from_grid(&self, grid: &[Complex<T>]) -> Vec<Complex<T>> {
        let mut result = Vec::with_capacity(self.n_nonuniform);

        for j in 0..self.n_nonuniform {
            let mut sum = Complex::<T>::zero();
            for &(idx, weight) in &self.spread_coeffs[j] {
                sum = sum + Complex::new(grid[idx].re * weight, grid[idx].im * weight);
            }
            result.push(sum);
        }

        result
    }

    /// Get the number of uniform grid points.
    pub fn n_uniform(&self) -> usize {
        self.n_uniform
    }

    /// Get the number of non-uniform points.
    pub fn n_nonuniform(&self) -> usize {
        self.n_nonuniform
    }

    /// Get the NUFFT type.
    pub fn nufft_type(&self) -> NufftType {
        self.nufft_type
    }

    /// Get the FFT planning flags used to build this plan's internal
    /// oversampled-grid FFT plans.
    pub fn flags(&self) -> Flags {
        self.flags
    }

    /// Get the normalized non-uniform points (in [0, 2π]).
    pub fn points(&self) -> &[f64] {
        &self.points
    }
}

/// Compute kernel width based on desired tolerance and oversampling ratio.
///
/// The spreading kernel is `exp(-β·(j/W)²)` with `β = 2.3·W` and
/// `W = kernel_width / 2`.  The required half-width W depends on both the
/// target accuracy and the oversampling ratio `σ`:
///
/// - Lower `σ` (e.g. 1.5) leaves fewer guard bands in the oversampled grid,
///   so the kernel must be wider to control sub-grid aliasing at Nyquist.
/// - Higher `σ` (e.g. 2.0) provides more isolation and needs fewer taps.
///
/// Empirical formula (validated against NUFFT tolerance sweep benchmarks):
///
/// ```text
/// W = ceil( -log10(tol) · (2 - σ/2) )
/// kw = max(4, 2·W)    (always even, so W = kw/2 is exact)
/// ```
///
/// Values:
/// - `σ=1.5, tol=1e-3`:  `W = ceil(3 · 1.25) = 4` → `kw = 8`
/// - `σ=2.0, tol=1e-3`:  `W = ceil(3 · 1.0)  = 3` → `kw = 6`
/// - `σ=1.5, tol=1e-6`:  `W = ceil(6 · 1.25) = 8` → `kw = 16`
/// - `σ=2.0, tol=1e-6`:  `W = ceil(6 · 1.0)  = 6` → `kw = 12`
///
/// The `default` parameter is the user-supplied `NufftOptions::kernel_width`
/// and sets an upper cap via `min(kw, default.max(12))`.
pub(crate) fn compute_kernel_width(tolerance: f64, oversampling: f64, default: usize) -> usize {
    let sigma = oversampling.clamp(1.05, 4.0); // guard against degenerate values
                                               // f(σ) = (2 - σ/2): factor that accounts for the reduction in guard-band
                                               // isolation as σ decreases toward 1.
    let f_sigma = 2.0 - sigma / 2.0;
    // W = half-width needed to achieve tolerance.
    let w = ((-tolerance.log10()) * f_sigma).ceil() as usize;
    // kw must be even (so that W = kw/2 exactly).
    // The lower bound is 4 (minimum useful kernel, W=2).
    // We take the max of the computed value and the user-supplied `default`
    // (from NufftOptions::kernel_width), so:
    //   - If the user requests a wider kernel, honour it.
    //   - If accuracy demands a wider kernel than `default`, use the wider one.
    (2 * w).max(4).max(default)
}

/// Map a uniform-grid output index `k` (in `0..n`) to the corresponding
/// oversampled-FFT-grid index and FFT-order deconvolution-table index, using
/// the **centered** frequency convention documented for [`Nufft::type1`] and
/// [`Nufft::type2`]:
///
/// ```text
/// freq = k - n/2
///   k=0     -> freq = -n/2   (most negative)
///   k=n/2   -> freq = 0      (DC)
///   k=n-1   -> freq = n/2-1  (most positive)
///
/// grid_idx   = freq                 if freq >= 0
///            = n_oversampled + freq if freq <  0
/// deconv_idx = freq                 if freq >= 0
///            = n + freq             if freq <  0
/// ```
///
/// `deconv_idx` indexes into the FFT-order table returned by
/// [`precompute_deconv_factors`]; `grid_idx` indexes into the FFT-order
/// oversampled-grid output.
///
/// This single convention is shared by the 1D, 2D ([`nufft2d`]), and 3D
/// ([`nufft3d`]) NUFFT implementations so that output index `k` always means
/// the same physical frequency regardless of dimensionality.
pub(crate) fn centered_freq_indices(k: usize, n: usize, n_oversampled: usize) -> (usize, usize) {
    let half_n = n / 2;
    let freq = (k as isize) - (half_n as isize);
    let grid_idx = if freq >= 0 {
        freq as usize
    } else {
        (n_oversampled as isize + freq) as usize
    };
    let deconv_idx = if freq >= 0 {
        freq as usize
    } else {
        (n as isize + freq) as usize
    };
    (grid_idx, deconv_idx)
}

/// Find next "smooth" number (product of small primes) for efficient FFT.
pub(crate) fn next_smooth_number(n: usize) -> usize {
    // Find next number that's a product of 2, 3, 5
    let mut candidate = n;
    loop {
        let mut temp = candidate;
        while temp.is_multiple_of(2) {
            temp /= 2;
        }
        while temp.is_multiple_of(3) {
            temp /= 3;
        }
        while temp.is_multiple_of(5) {
            temp /= 5;
        }
        if temp == 1 {
            return candidate;
        }
        candidate += 1;
    }
}

/// Precompute spreading coefficients using Gaussian kernel.
pub(crate) fn precompute_spreading_coeffs<T: Float>(
    points: &[f64],
    n_grid: usize,
    kernel_width: usize,
) -> Vec<Vec<(usize, T)>> {
    let grid_spacing = 2.0 * core::f64::consts::PI / (n_grid as f64);
    let half_width = kernel_width / 2;

    // Gaussian kernel parameter (beta).
    //
    // The kernel is exp(-β · (dx / (h · W))²) with W = half_width.
    // β must scale with W (not kernel_width) so that the edge weight
    // exp(-β) remains finite regardless of how large kw gets.  Using
    // β = 2.3 · W gives ~exp(-2.3) ≈ 0.10 edge weight, which provides
    // enough taper while keeping sub-grid variation manageable.
    let beta = 2.3 * (half_width as f64);

    points
        .iter()
        .map(|&x| {
            // Find nearest grid point
            let grid_pos = x / grid_spacing;
            let center = grid_pos.round() as isize;

            let mut coeffs = Vec::with_capacity(kernel_width);

            for offset in -(half_width as isize)..=(half_width as isize) {
                let grid_idx = (center + offset).rem_euclid(n_grid as isize) as usize;
                let grid_x = (grid_idx as f64) * grid_spacing;

                // Distance from point to grid location
                let mut dx = x - grid_x;
                // Wrap around
                if dx > core::f64::consts::PI {
                    dx -= 2.0 * core::f64::consts::PI;
                } else if dx < -core::f64::consts::PI {
                    dx += 2.0 * core::f64::consts::PI;
                }

                // Gaussian kernel: exp(-beta * (dx/width)^2)
                let normalized_dx = dx / (grid_spacing * (half_width as f64));
                let weight = (-beta * normalized_dx * normalized_dx).exp();

                if weight > 1e-15 {
                    coeffs.push((grid_idx, T::from_f64(weight)));
                }
            }

            coeffs
        })
        .collect()
}

/// Precompute deconvolution factors for the Gaussian NUFFT.
///
/// The spreading kernel in [`precompute_spreading_coeffs`] is:
/// ```text
/// w(dx) = exp(-β · (dx / (h · W_int))²)
///   h     = 2π / n_oversampled   (oversampled grid spacing)
///   W_int = kernel_width / 2     (integer half-width, support is 2·W_int+1 points)
///   β     = 2.3 · W_int   (= 2.3 · kernel_width / 2)
/// ```
///
/// ## Exact discrete kernel DFT
///
/// The kernel is symmetric about its centre, so its DFT (evaluated at the
/// integer output frequency `freq`) is real and equals:
/// ```text
/// ψ̂(freq) = Σ_{j=-W_int}^{W_int}  exp(-β · (j/W_int)²) · cos(2π · freq · j / n_oversampled)
/// ```
/// The deconvolution factor is `1 / ψ̂(freq)`.
///
/// ## Phase correction for the [0, 2π] shift
///
/// Points are normalised from `[-π, π]` to `[0, 2π]` by adding π.  This shift
/// multiplies the DFT result at oversampled bin `k` by `exp(-i·k·π) = (-1)^k`.
/// After deconvolution the residual phase `(-1)^k` must be removed by
/// multiplying by `exp(+i·k·π) = (-1)^k`.  Because `(-1)^k = ±1` the combined
/// factor remains real-valued.
///
/// The oversampled grid bin `k` that holds frequency `freq` is:
/// - `k = freq`              when `freq ≥ 0`
/// - `k = n_oversampled + freq`  when `freq < 0`
///
/// In FFT-order deconvolution index `d`:
/// - `d < N/2`: `freq = d`,  `k = d`
/// - `d ≥ N/2`: `freq = d − N`, `k = n_oversampled + d − N`
///
/// The returned factors are indexed in **FFT order** (`d = 0..N-1`).
pub(crate) fn precompute_deconv_factors<T: Float>(
    n_uniform: usize,
    n_oversampled: usize,
    kernel_width: usize,
) -> Vec<Complex<T>> {
    let w_int = (kernel_width / 2) as isize; // integer half-width for spreading
                                             // β must match the β used in precompute_spreading_coeffs: β = 2.3 · W_int
    let beta = 2.3 * (w_int as f64);
    let two_pi_over_nos = 2.0 * core::f64::consts::PI / (n_oversampled as f64);

    (0..n_uniform)
        .map(|d| {
            // FFT-order frequency and oversampled grid bin
            let (freq, grid_bin) = if d < n_uniform / 2 {
                (d as isize, d)
            } else {
                let f = (d as isize) - (n_uniform as isize);
                (f, n_oversampled + d - n_uniform)
            };

            // Exact kernel DFT at `freq` (cosine sum over kernel support):
            //   ψ̂(freq) = Σ_{j=-W_int}^{W_int} exp(-β·(j/W_int)²) · cos(2π·freq·j/n_os)
            let kernel_dft: f64 = (-w_int..=w_int)
                .map(|j| {
                    let w = (-beta * ((j * j) as f64 / (w_int * w_int) as f64)).exp();
                    let angle = two_pi_over_nos * (freq * j) as f64;
                    w * angle.cos()
                })
                .sum();

            // Phase correction: spreading used x_norm = x_orig + π, which added
            // exp(-i·k·π) = (-1)^k to every DFT bin.  Multiply by (-1)^k to undo.
            let phase_sign = if grid_bin % 2 == 0 { 1.0_f64 } else { -1.0_f64 };

            // Deconvolution: 1/ψ̂(freq) with phase correction (always real).
            let deconv = if kernel_dft.abs() > f64::EPSILON {
                phase_sign / kernel_dft
            } else {
                0.0_f64
            };

            Complex::new(T::from_f64(deconv), T::ZERO)
        })
        .collect()
}

// Convenience functions

/// Compute Type 1 NUFFT (non-uniform → uniform).
///
/// # Arguments
///
/// * `points` - Non-uniform sample locations in [-π, π]
/// * `values` - Complex values at the sample points
/// * `n_output` - Number of uniform output frequencies
/// * `tolerance` - Target relative accuracy (e.g., 1e-6)
///
/// # Returns
///
/// Uniform Fourier coefficients.
///
/// # Errors
///
/// Returns error if points are out of range or tolerance is invalid.
pub fn nufft_type1<T: Float>(
    points: &[f64],
    values: &[Complex<T>],
    n_output: usize,
    tolerance: f64,
) -> NufftResult<Vec<Complex<T>>> {
    let plan = Nufft::new(NufftType::Type1, n_output, points, tolerance)?;
    plan.type1(values)
}

/// Compute Type 2 NUFFT (uniform → non-uniform).
///
/// # Arguments
///
/// * `coeffs` - Uniform Fourier coefficients
/// * `points` - Non-uniform evaluation points in [-π, π]
/// * `tolerance` - Target relative accuracy (e.g., 1e-6)
///
/// # Returns
///
/// Values at non-uniform points.
///
/// # Errors
///
/// Returns error if points are out of range or tolerance is invalid.
pub fn nufft_type2<T: Float>(
    coeffs: &[Complex<T>],
    points: &[f64],
    tolerance: f64,
) -> NufftResult<Vec<Complex<T>>> {
    let plan = Nufft::new(NufftType::Type2, coeffs.len(), points, tolerance)?;
    plan.type2(coeffs)
}

/// Compute Type 3 NUFFT (non-uniform → non-uniform).
///
/// # Arguments
///
/// * `source_points` - Source non-uniform locations in [-π, π]
/// * `values` - Complex values at source points
/// * `target_points` - Target non-uniform locations in [-π, π]
/// * `tolerance` - Target relative accuracy
///
/// # Returns
///
/// Values at target points.
///
/// # Errors
///
/// Returns error if points are out of range or tolerance is invalid.
pub fn nufft_type3<T: Float>(
    source_points: &[f64],
    values: &[Complex<T>],
    target_points: &[f64],
    tolerance: f64,
) -> NufftResult<Vec<Complex<T>>> {
    // Use intermediate uniform grid size based on source and target counts
    let n_uniform = (source_points.len() + target_points.len()).next_power_of_two();
    let plan = Nufft::new(NufftType::Type1, n_uniform, source_points, tolerance)?;
    plan.execute_type3(values, target_points)
}

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

    // -----------------------------------------------------------------------
    // Dense NDFT reference helpers, using the centered frequency convention
    // (freq = k - n/2) documented on [`Nufft::type1`] / [`Nufft::type2`].
    // A more extensive tolerance x oversampling x size sweep against the same
    // reference lives in `oxifft/tests/nufft_tolerance_sweep.rs`; these local
    // helpers keep the module's own unit tests self-contained and numeric.
    // -----------------------------------------------------------------------

    /// Dense NDFT Type 1 (non-uniform -> uniform), O(n*m).
    fn dense_ndft_type1(points: &[f64], values: &[Complex<f64>], n: usize) -> Vec<Complex<f64>> {
        let half = (n / 2) as isize;
        (0..n)
            .map(|k| {
                let freq = (k as isize - half) as f64;
                values
                    .iter()
                    .zip(points.iter())
                    .fold(Complex::new(0.0, 0.0), |acc, (&val, &xj)| {
                        let angle = -freq * xj;
                        acc + val * Complex::new(angle.cos(), angle.sin())
                    })
            })
            .collect()
    }

    /// Dense NDFT Type 2 (uniform -> non-uniform), O(n*m).
    fn dense_ndft_type2(fhat: &[Complex<f64>], points: &[f64]) -> Vec<Complex<f64>> {
        let n = fhat.len();
        let half = (n / 2) as isize;
        points
            .iter()
            .map(|&xj| {
                (0..n).fold(Complex::new(0.0, 0.0), |acc, k| {
                    let freq = (k as isize - half) as f64;
                    let angle = freq * xj;
                    acc + fhat[k] * Complex::new(angle.cos(), angle.sin())
                })
            })
            .collect()
    }

    /// Maximum `|nufft[i] - ref[i]| / max(|ref[i]|)` over all bins.
    fn max_relative_error(nufft_out: &[Complex<f64>], reference: &[Complex<f64>]) -> f64 {
        let ref_max = reference.iter().map(|c| c.norm()).fold(0.0_f64, f64::max);
        if ref_max < 1e-30 {
            return 0.0;
        }
        nufft_out
            .iter()
            .zip(reference.iter())
            .map(|(n, r)| (*n - *r).norm() / ref_max)
            .fold(0.0_f64, f64::max)
    }

    /// Relative-error headroom for tolerance 1e-6 (10x margin, matching the
    /// dedicated tolerance sweep in `oxifft/tests/nufft_tolerance_sweep.rs`).
    const HEADROOM: f64 = 1e-5;

    #[test]
    fn test_nufft_type1_uniform_points_matches_dense_ndft() {
        // If points are uniformly spaced, NUFFT Type 1 should closely match
        // the dense NDFT (and, by construction, the regular DFT).
        let n = 8;
        let points: Vec<f64> = (0..n)
            .map(|k| -core::f64::consts::PI + (k as f64) * 2.0 * core::f64::consts::PI / (n as f64))
            .collect();

        let values: Vec<Complex<f64>> = (0..n)
            .map(|k| Complex::new((k as f64).cos(), (k as f64).sin()))
            .collect();

        let result = nufft_type1(&points, &values, n, 1e-6).expect("NUFFT Type1 failed");
        assert_eq!(result.len(), n);

        let reference = dense_ndft_type1(&points, &values, n);
        let rel_err = max_relative_error(&result, &reference);
        assert!(
            rel_err <= HEADROOM,
            "uniform-points Type1 rel_err {rel_err:.2e} exceeds headroom {HEADROOM:.2e}"
        );
    }

    #[test]
    fn test_nufft_type2_single_frequency_matches_dense_ndft() {
        // Single frequency coefficient should produce a pure sinusoid at
        // frequency `freq = k - n/2` at each evaluation point.
        let n = 16;
        let mut coeffs = vec![Complex::<f64>::zero(); n];
        coeffs[1] = Complex::new(1.0, 0.0); // freq = 1 - n/2 = -7

        let points: Vec<f64> = (0..5)
            .map(|k| -core::f64::consts::PI + f64::from(k) * 0.5)
            .collect();

        let result = nufft_type2(&coeffs, &points, 1e-6).expect("NUFFT Type2 failed");
        assert_eq!(result.len(), 5);

        let reference = dense_ndft_type2(&coeffs, &points);
        let rel_err = max_relative_error(&result, &reference);
        assert!(
            rel_err <= HEADROOM,
            "single-frequency Type2 rel_err {rel_err:.2e} exceeds headroom {HEADROOM:.2e}"
        );
    }

    #[test]
    fn test_nufft_type1_and_type2_independently_match_dense_ndft() {
        // Type 1 and Type 2 are each validated independently against their
        // own dense NDFT reference (Type1-then-Type2 is not an exact
        // identity in general since Type 1 is the adjoint, not the inverse,
        // of Type 2).
        let n = 32;
        let points: Vec<f64> = (0..10).map(|k| -2.5 + f64::from(k) * 0.5).collect();

        let values: Vec<Complex<f64>> = points
            .iter()
            .map(|&x| Complex::new(x.cos(), x.sin()))
            .collect();

        let type1_out = nufft_type1(&points, &values, n, 1e-6).expect("Type1 failed");
        let type1_ref = dense_ndft_type1(&points, &values, n);
        let type1_rel_err = max_relative_error(&type1_out, &type1_ref);
        assert!(
            type1_rel_err <= HEADROOM,
            "Type1 rel_err {type1_rel_err:.2e} exceeds headroom {HEADROOM:.2e}"
        );

        let type2_out = nufft_type2(&type1_ref, &points, 1e-6).expect("Type2 failed");
        let type2_ref = dense_ndft_type2(&type1_ref, &points);
        let type2_rel_err = max_relative_error(&type2_out, &type2_ref);
        assert!(
            type2_rel_err <= HEADROOM,
            "Type2 rel_err {type2_rel_err:.2e} exceeds headroom {HEADROOM:.2e}"
        );
        assert_eq!(type2_out.len(), values.len());
    }

    #[test]
    fn test_nufft_type3_matches_composed_dense_reference() {
        // `execute_type3` is defined as Type1 (to an intermediate uniform
        // grid of `n_uniform` bandlimited modes) followed by Type2 at the
        // target points, so its ground truth is the *composition* of the two
        // dense NDFT references (not a "true" continuous non-uniform-to-
        // non-uniform transform), matching the documented algorithm.
        let n_uniform = 32;
        let source_points: Vec<f64> = (0..9).map(|k| -2.8 + f64::from(k) * 0.6).collect();
        let target_points: Vec<f64> = (0..6).map(|k| -1.9 + f64::from(k) * 0.55).collect();

        let values: Vec<Complex<f64>> = source_points
            .iter()
            .map(|&x| Complex::new((x * 0.5).cos(), (x * 0.5).sin()))
            .collect();

        let plan = Nufft::new(NufftType::Type1, n_uniform, &source_points, 1e-6)
            .expect("Type3 plan failed");
        let result = plan
            .execute_type3(&values, &target_points)
            .expect("Type3 execute failed");
        assert_eq!(result.len(), target_points.len());

        let intermediate_ref = dense_ndft_type1(&source_points, &values, n_uniform);
        let composed_ref = dense_ndft_type2(&intermediate_ref, &target_points);

        let rel_err = max_relative_error(&result, &composed_ref);
        assert!(
            rel_err <= HEADROOM,
            "Type3 rel_err {rel_err:.2e} exceeds headroom {HEADROOM:.2e}"
        );
    }

    /// Coincident non-uniform points must combine linearly: two points at
    /// the same location with values `c1`,`c2` must give exactly the same
    /// Type 1 result (up to floating rounding) as a single point at that
    /// location with value `c1+c2`.
    #[test]
    fn test_nufft_type1_coincident_points_linearity() {
        let n = 16;
        let x_single = vec![0.4f64];
        let c_combined = vec![Complex::new(1.7, -0.5)];

        let x_dup = vec![0.4f64, 0.4f64];
        let c_dup = vec![Complex::new(1.0, -0.2), Complex::new(0.7, -0.3)];

        let result_single = nufft_type1(&x_single, &c_combined, n, 1e-6).expect("single failed");
        let result_dup = nufft_type1(&x_dup, &c_dup, n, 1e-6).expect("dup failed");

        for (a, b) in result_single.iter().zip(result_dup.iter()) {
            assert!(
                (*a - *b).norm() < 1e-9,
                "coincident-point linearity violated: {a:?} vs {b:?}"
            );
        }
    }

    /// Points exactly at the domain edges (`-π`, `π`) must not panic and
    /// must still match the dense reference.
    #[test]
    fn test_nufft_type1_edge_points_matches_dense_ndft() {
        let pi = core::f64::consts::PI;
        let points = vec![-pi, pi, 0.0];
        let values = vec![
            Complex::new(1.0, 0.0),
            Complex::new(0.5, -0.5),
            Complex::new(-0.3, 0.2),
        ];
        let n = 16;

        let result = nufft_type1(&points, &values, n, 1e-6).expect("edge-point Type1 failed");
        let reference = dense_ndft_type1(&points, &values, n);
        let rel_err = max_relative_error(&result, &reference);
        assert!(
            rel_err <= HEADROOM,
            "edge-points rel_err {rel_err:.2e} exceeds headroom {HEADROOM:.2e}"
        );
    }

    /// Empty non-uniform input must yield an all-zero uniform grid (not a
    /// panic or garbage output).
    #[test]
    fn test_nufft_type1_empty_input_returns_zero_grid() {
        let points: Vec<f64> = vec![];
        let values: Vec<Complex<f64>> = vec![];
        let n = 8;

        let result = nufft_type1(&points, &values, n, 1e-6).expect("empty Type1 failed");
        assert_eq!(result.len(), n);
        for v in &result {
            assert_eq!(v.re, 0.0);
            assert_eq!(v.im, 0.0);
        }
    }

    /// Empty evaluation-point input must yield an empty result (not a panic).
    #[test]
    fn test_nufft_type2_empty_points_returns_empty() {
        let n = 8;
        let coeffs = vec![Complex::<f64>::zero(); n];
        let points: Vec<f64> = vec![];

        let result = nufft_type2(&coeffs, &points, 1e-6).expect("empty Type2 failed");
        assert!(result.is_empty());
    }

    #[test]
    fn test_nufft_error_handling() {
        let points = vec![0.0, 0.5, 1.0];

        // Invalid size
        let result = Nufft::<f64>::new(NufftType::Type1, 0, &points, 1e-6);
        assert!(result.is_err());

        // Point out of range
        let bad_points = vec![0.0, 5.0]; // 5.0 > π
        let result = Nufft::<f64>::new(NufftType::Type1, 16, &bad_points, 1e-6);
        assert!(result.is_err());

        // Invalid tolerance
        let result = Nufft::<f64>::new(NufftType::Type1, 16, &points, -1e-6);
        assert!(result.is_err());
    }

    #[test]
    fn test_smooth_number() {
        assert_eq!(next_smooth_number(100), 100); // 100 = 2^2 * 5^2
        assert_eq!(next_smooth_number(101), 108); // 108 = 2^2 * 3^3
        assert_eq!(next_smooth_number(7), 8); // 8 = 2^3
    }

    #[test]
    fn test_nufft_new_defaults_to_estimate_flag() {
        // `new`/`with_options` must no longer hardcode Flags::MEASURE.
        let points = vec![0.0, 0.3, -0.5];
        let plan = Nufft::<f64>::new(NufftType::Type1, 16, &points, 1e-6).expect("plan failed");
        assert_eq!(plan.flags(), Flags::ESTIMATE);
    }

    #[test]
    fn test_nufft_with_flags_estimate_and_measure_match() {
        let n = 16;
        let points: Vec<f64> = (0..6).map(|k| -2.0 + f64::from(k) * 0.6).collect();
        let values: Vec<Complex<f64>> = points
            .iter()
            .map(|&x| Complex::new(x.cos(), x.sin()))
            .collect();

        let plan_estimate =
            Nufft::<f64>::with_flags(NufftType::Type1, n, &points, 1e-6, Flags::ESTIMATE)
                .expect("ESTIMATE plan failed");
        let plan_measure =
            Nufft::<f64>::with_flags(NufftType::Type1, n, &points, 1e-6, Flags::MEASURE)
                .expect("MEASURE plan failed");

        assert_eq!(plan_estimate.flags(), Flags::ESTIMATE);
        assert_eq!(plan_measure.flags(), Flags::MEASURE);

        let out_estimate = plan_estimate
            .type1(&values)
            .expect("ESTIMATE execute failed");
        let out_measure = plan_measure.type1(&values).expect("MEASURE execute failed");

        for (a, b) in out_estimate.iter().zip(out_measure.iter()) {
            assert!((*a - *b).norm() < 1e-9);
        }
    }

    #[test]
    fn test_nufft_type2_forward_and_inverse_plans_share_flags() {
        // Regression test for the forward=MEASURE/inverse=ESTIMATE flag
        // inconsistency: both the forward (Type 1 spreading) and inverse
        // (Type 2 interpolation) FFT plans must be built with the same,
        // caller-chosen flags, and the inverse plan must be precomputed at
        // construction time (not re-planned on every `type2()` call).
        let n = 16;
        let coeffs: Vec<Complex<f64>> = (0..n)
            .map(|k| Complex::new(((k as f64) * 0.1).cos(), 0.0))
            .collect();
        let points: Vec<f64> = (0..5).map(|k| -2.0 + f64::from(k) * 0.7).collect();

        let plan = Nufft::<f64>::with_flags(NufftType::Type2, n, &points, 1e-6, Flags::MEASURE)
            .expect("plan failed");
        assert_eq!(plan.flags(), Flags::MEASURE);

        // Calling type2() repeatedly must succeed and give consistent output
        // now that the inverse plan is cached rather than rebuilt each call.
        let out1 = plan.type2(&coeffs).expect("first type2 call failed");
        let out2 = plan.type2(&coeffs).expect("second type2 call failed");
        assert_eq!(out1.len(), out2.len());
        for (a, b) in out1.iter().zip(out2.iter()) {
            assert_eq!(a.re, b.re);
            assert_eq!(a.im, b.im);
        }
    }
}