rexafs 0.2.5

// Standard library dependencies

// External dependencies
use easyfft::prelude::DynRealFft;
use easyfft::{dyn_size::realfft::DynRealDft, num_complex::Complex};
use nalgebra::{DVector, Owned};
use ndarray::{
    Array, Array1, ArrayBase, ArrayView, ArrayView1, Axis, Ix, Ix1, OwnedRepr, ViewRepr,
};
use num_complex::Complex64;
use serde::{Deserialize, Serialize};

// load dependencies
use super::errors::FFTError;
pub use super::fft_grid::FFTGrid;

// Load local traits
use super::mathutils::MathUtils;
use super::xafsutils::ftwindow;
use crate::xafs::xafsutils::FTWindow;

/// Forward EXAFS transform from chi(k) to complex chi(R).
///
/// The transform is an unnormalized negative-exponent FFT multiplied by
/// `kstep / sqrt(pi)`; no additional `1/nfft`, phase factor, or window-area
/// normalization is applied. Use `new()` for recommended settings, then `xftf`
/// to calculate. Mutating these standalone fields does not recalculate caches;
/// use `Spectrum::set_fft` when managing a spectrum pipeline.
///
/// The sign and FFT normalization match the
/// [NumPy convention](https://numpy.org/doc/stable/reference/routines.fft.html#normalization).
///
/// On the zero-origin uniform grid k_j=j*dk, let
/// g_j=chi(k_j)*k_j^w*W(k_j), with w the integer k-weight and W the window.
/// For N=`nfft`, dk=`kstep` in Å⁻¹ and i²=−1, the implementation is
///
/// ```text
/// chi_R[m] = dk/sqrt(pi) * sum(j=0..N-1, g_j * exp(-2*pi*i*j*m/N))
/// R[m]     = pi*m/(N*dk)
/// ```
///
/// Zero padding supplies g_j=0 beyond the prepared data. R is in Å, and
/// chi_R has units Å^(-(w+1)) when chi is dimensionless. This sum approximates
/// the convention `1/sqrt(pi) * integral(chi(k)*k^w*W(k)*exp(-2*i*k*R), dk)`.
/// The kernel's 2*k*R is the EXAFS path oscillation; scattering phase means an
/// uncorrected peak R is not directly a bond length. See
/// [Rehr and Albers (2000)](https://doi.org/10.1103/RevModPhys.72.621).
/// Resolved automatic parameters are retained in this object, not recomputed
/// from each new input; clear the relevant field to None or use fresh settings.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default)]
pub struct XrayFFTF {
    /// Sampling and window domain. Default: `FFTGrid::Larch` in this ndarray compatibility backend.
    /// Input preserves samples; Larch resamples on a uniform grid starting at zero.
    pub grid: FFTGrid,
    /// Maximum displayed R in Å; default 10.0. The full transform is retained for
    /// inverse filtering, so this does not set an R-space filter cutoff.
    pub rmax_out: Option<f64>,
    /// Window family; default `Some(FTWindow::KaiserBessel)`. `None` selects
    /// the window helper's Hanning default, not an all-pass window.
    pub window: Option<FTWindow>,
    /// Low-side window parameter; default 1.0. Usually a taper width in Å⁻¹;
    /// Kaiser–Bessel also uses its numerical value as a shape parameter. Gaussian
    /// uses it as the width, and fractional Hanning uses fractional taper geometry.
    pub dk: Option<f64>,
    /// High-side window parameter; `None` uses `dk`. Interpretation depends on
    /// window family; equal values do not make different families equivalent.
    pub dk2: Option<f64>,
    /// Lower window boundary in Å⁻¹; default 2.0. `None` uses the first input k.
    pub kmin: Option<f64>,
    /// Upper window boundary in Å⁻¹; default 15.0. `None` uses the last input k.
    /// In Larch mode, the window domain extends through `kmax + dk2`.
    pub kmax: Option<f64>,
    /// Power of k multiplying chi before transformation; default 2.0.
    /// Nonnegative fractional values are floored. Larger values emphasize high-k
    /// signal and noise; this weighting is separate from AUTOBK's own k-weight.
    pub kweight: Option<f64>,
    /// FFT length; default 2048. Short input is zero-padded; Input mode truncates
    /// longer input to this length. Larch mode requires room for its extended window.
    /// More zero-padding refines the R grid without adding experimental resolution.
    pub nfft: Option<usize>,
    /// FFT sample spacing in Å⁻¹; `None` infers `k[1] - k[0]`. This controls both
    /// R spacing `pi / (nfft * kstep)` and amplitude `kstep / sqrt(pi)`.
    /// AUTOBK normally supplies 0.05 Å⁻¹; Input mode does not resample when this changes.
    pub kstep: Option<f64>,
    /// Displayed R grid in Å, populated by `xftf`; `None` before calculation.
    pub r: Option<ArrayBase<OwnedRepr<f64>, Ix1>>,
    // currently asking for serde support in the easyfft crate
    /// Full complex real-FFT coefficients, including DC and the even-length Nyquist
    /// bin. For dimensionless chi and k-weight w, units are Å⁻⁽ʷ⁺¹⁾. This cache is
    /// not truncated by `rmax_out`; `None` before calculation.
    pub chir: Option<DynRealDft<f64>>,
    /// Magnitude of the displayed complex transform, paired with `r`.
    /// Units are Å⁻⁽ʷ⁺¹⁾ for dimensionless chi; `None` before calculation.
    pub chir_mag: Option<ArrayBase<OwnedRepr<f64>, Ix1>>,
    /// Dimensionless window on the prepared k grid, before k-weight multiplication.
    /// Use `Spectrum::kwin_k` to obtain the matching axis; `None` before calculation.
    pub kwin: Option<ArrayBase<OwnedRepr<f64>, Ix1>>,
}

// Preserve the original equality contract: the complex FFT cache is omitted.
impl PartialEq for XrayFFTF {
    fn eq(&self, other: &Self) -> bool {
        self.grid == other.grid
            && self.rmax_out == other.rmax_out
            && self.window == other.window
            && self.dk == other.dk
            && self.dk2 == other.dk2
            && self.kmin == other.kmin
            && self.kmax == other.kmax
            && self.kweight == other.kweight
            && self.nfft == other.nfft
            && self.kstep == other.kstep
            && self.r == other.r
            && self.chir_mag == other.chir_mag
            && self.kwin == other.kwin
    }
}

impl Default for XrayFFTF {
    fn default() -> Self {
        XrayFFTF {
            grid: FFTGrid::Larch,
            rmax_out: Some(10.0),
            window: Some(FTWindow::KaiserBessel),
            dk: Some(1.),
            dk2: None,
            kmin: Some(2.0),
            kmax: Some(15.0),
            kweight: Some(2.0),
            nfft: Some(2048),
            kstep: None,
            r: None,
            chir: None,
            chir_mag: None,
            kwin: None,
        }
    }
}

impl XrayFFTF {
    /// Create recommended forward settings with automatic input spacing.
    pub fn new() -> XrayFFTF {
        XrayFFTF::default()
    }

    /// Resolve automatic settings from at least two k samples without calculating.
    ///
    /// Prefer the validated transform method for public input. This lower-level
    /// helper indexes the supplied grid directly and can panic on an empty grid.
    /// Only None fields are resolved; their resulting values remain stored when
    /// later inputs change. Reset them to None to infer from a new grid.
    pub fn fill_parameter(&mut self, k: ArrayBase<ViewRepr<&f64>, Ix1>) -> &mut Self {
        if self.kweight.is_none() {
            self.kweight = Some(2.0);
        }

        self.kweight = Some(self.kweight.unwrap().max(0.0).floor());

        if self.kstep.is_none() {
            self.kstep = Some(k[1] - k[0]);
        }

        if self.kmin.is_none() {
            self.kmin = Some(k[0]);
        }

        if self.kmax.is_none() {
            self.kmax = Some(k[k.len() - 1]);
        }

        if self.dk.is_none() {
            self.dk = Some(1.0);
        }

        if self.dk2.is_none() {
            self.dk2 = self.dk;
        }

        if self.nfft.is_none() {
            self.nfft = Some(2048);
        }

        if self.rmax_out.is_none() {
            self.rmax_out = Some(10.0);
        }

        self
    }

    /// Validate input/settings, resolve automatic values and allocate prepared data.
    /// Returns `(chi * k.powi(kweight), window)` before multiplying the window.
    /// The Larch default resamples from k=0; Input preserves the supplied samples.
    /// No Fourier coefficients are calculated by this preparation step.
    pub fn xftf_prep(
        &mut self,
        k: ArrayBase<ViewRepr<&f64>, Ix1>,
        chi: ArrayBase<ViewRepr<&f64>, Ix1>,
    ) -> Result<
        (
            ArrayBase<OwnedRepr<f64>, Ix1>,
            ArrayBase<OwnedRepr<f64>, Ix1>,
        ),
        FFTError,
    > {
        if self.nfft.is_some_and(|n| n < 2) {
            return Err(FFTError::InvalidParameter {
                parameter: "nfft".into(),
                reason: "must be at least 2".into(),
            });
        }
        for (name, value, strictly_positive) in [
            ("kstep", self.kstep, true),
            ("kweight", self.kweight, false),
            ("dk", self.dk, false),
            ("dk2", self.dk2, false),
            ("rmax_out", self.rmax_out, false),
        ] {
            if value.is_some_and(|v| !v.is_finite() || v < 0.0 || (strictly_positive && v == 0.0)) {
                return Err(FFTError::InvalidParameter {
                    parameter: name.into(),
                    reason: "must be finite and nonnegative (kstep must be positive)".into(),
                });
            }
        }
        if self.kmin.is_some_and(|v| !v.is_finite())
            || self.kmax.is_some_and(|v| !v.is_finite())
            || self.kmin.zip(self.kmax).is_some_and(|(lo, hi)| lo >= hi)
        {
            return Err(FFTError::InvalidParameter {
                parameter: "kmin/kmax".into(),
                reason: "must be finite with kmin < kmax".into(),
            });
        }
        if k.len() != chi.len() {
            return Err(FFTError::InterpolationFailed {
                reason: "k/chi length mismatch".to_string(),
            });
        }
        if k.len() < 2 {
            let (kmin, kmax) = if k.is_empty() {
                (0.0, 0.0)
            } else {
                (k[0], k[k.len() - 1])
            };
            return Err(FFTError::InsufficientPoints {
                min: 2,
                actual: k.len(),
                kmin,
                kmax,
            });
        }

        if k.iter().chain(chi.iter()).any(|v| !v.is_finite())
            || k.iter().zip(k.iter().skip(1)).any(|(a, b)| a >= b)
        {
            return Err(FFTError::InvalidParameter {
                parameter: "k/chi".into(),
                reason: "must be finite with strictly increasing k".into(),
            });
        }
        self.fill_parameter(k);
        let kweight = self.kweight.unwrap() as i32;
        if self.grid == FFTGrid::Input {
            let win = self
                .window
                .unwrap_or_default()
                .window(&k.to_owned(), self.kmin, self.kmax, self.dk, self.dk2)
                .map_err(|e| FFTError::WindowCalculationFailed {
                    reason: e.to_string(),
                })?;
            return Ok((chi.to_owned() * k.mapv(|v| v.powi(kweight)), win));
        }
        let step = self.kstep.unwrap();
        let last = k[k.len() - 1];
        let npts = (1.01 + last / step).floor();
        let extent = last.max(self.kmax.unwrap() + self.dk2.unwrap());
        let nwin = (1.01 + extent / step).floor();
        if k[0] < 0.0 || !nwin.is_finite() || npts < 2.0 || nwin > self.nfft.unwrap() as f64 {
            return Err(FFTError::InvalidParameter {
                parameter: "Larch grid".into(),
                reason: "requires nonnegative k, at least two resampled points, and nfft large enough for kmax+dk2".into(),
            });
        }
        let npts = npts as usize;
        let k_ = Array1::from_iter((0..nwin as usize).map(|i| i as f64 * step));

        let chi_ = if let (Some(k_slice), Some(chi_slice)) = (k.as_slice(), chi.as_slice()) {
            k_.interpolate(k_slice, chi_slice)
        } else {
            k_.interpolate(&k.to_vec(), &chi.to_vec())
        }
        .map_err(|e| FFTError::InterpolationFailed {
            reason: e.to_string(),
        })?;
        let win = self
            .window
            .unwrap_or_default()
            .window(&k_, self.kmin, self.kmax, self.dk, self.dk2)
            .map_err(|e| FFTError::WindowCalculationFailed {
                reason: e.to_string(),
            })?;
        let win = (win).slice_axis(Axis(0), (0..npts).into()).to_owned();
        let chi_ = &chi_.slice_axis(Axis(0), (0..npts).into())
            * &k_
                .slice_axis(Axis(0), (0..npts).into())
                .map(|x| x.powi(kweight));

        Ok((chi_, win))
    }

    /// Compute the weighted, windowed forward transform and replace cached results.
    ///
    /// Requires matched finite arrays with at least two increasing k samples. The
    /// compatibility default is Larch resampling from zero. Invalid settings or
    /// an FFT too short for that extended grid return `FFTError`.
    pub fn xftf(
        &mut self,
        k: ArrayBase<ViewRepr<&f64>, Ix1>,
        chi: ArrayBase<ViewRepr<&f64>, Ix1>,
    ) -> Result<&mut Self, FFTError> {
        let (cchi, win) = self.xftf_prep(k, chi)?;

        let cchi_fft = xftf_fast(
            (&cchi * &win).view(),
            self.nfft.unwrap(),
            self.kstep.unwrap(),
        );

        let rstep = std::f64::consts::PI / self.kstep.unwrap() / self.nfft.unwrap() as f64;

        // The length of r is different by 1 between xraylarch and rexafs. This is due to the implementation of FFT.
        let irmax =
            (self.nfft.unwrap() / 2 + 1).min((1.01 + self.rmax_out.unwrap() / rstep) as usize);

        self.r = Some(Array1::range(0.0, irmax as f64 * rstep, rstep));

        self.chir = Some(cchi_fft.clone());
        self.chir_mag = Some(cchi_fft[0..irmax].norm());
        self.kwin = Some(win);

        Ok(self)
    }

    /// Read the stored value without recalculating.
    ///
    /// Maximum displayed R in Å; default 10.0. The full transform is retained for
    /// inverse filtering, so this does not set an R-space filter cutoff.
    pub fn get_rmax_out(&self) -> Option<&f64> {
        self.rmax_out.as_ref()
    }

    /// Read the stored value without recalculating.
    ///
    /// Window family; default `Some(FTWindow::KaiserBessel)`. `None` selects
    /// the window helper's Hanning default, not an all-pass window.
    pub fn get_window(&self) -> Option<&FTWindow> {
        self.window.as_ref()
    }

    /// Read the stored value without recalculating.
    ///
    /// Low-side window parameter; default 1.0. Usually a taper width in Å⁻¹;
    /// Kaiser–Bessel also uses its numerical value as a shape parameter. Gaussian
    /// uses it as the width, and fractional Hanning uses fractional taper geometry.
    pub fn get_dk(&self) -> Option<&f64> {
        self.dk.as_ref()
    }

    /// Read the stored value without recalculating.
    ///
    /// High-side window parameter; `None` uses `dk`. Interpretation depends on
    /// window family; equal values do not make different families equivalent.
    pub fn get_dk2(&self) -> Option<&f64> {
        self.dk2.as_ref()
    }

    /// Read the stored value without recalculating.
    ///
    /// Lower window boundary in Å⁻¹; default 2.0. `None` uses the first input k.
    pub fn get_kmin(&self) -> Option<&f64> {
        self.kmin.as_ref()
    }

    /// Read the stored value without recalculating.
    ///
    /// Upper window boundary in Å⁻¹; default 15.0. `None` uses the last input k.
    /// In Larch mode, the window domain extends through `kmax + dk2`.
    pub fn get_kmax(&self) -> Option<&f64> {
        self.kmax.as_ref()
    }

    /// Read the stored value without recalculating.
    ///
    /// Power of k multiplying chi before transformation; default 2.0.
    /// Nonnegative fractional values are floored. Larger values emphasize high-k
    /// signal and noise; this weighting is separate from AUTOBK's own k-weight.
    pub fn get_kweight(&self) -> Option<&f64> {
        self.kweight.as_ref()
    }

    /// Read the stored value without recalculating.
    ///
    /// Displayed R grid in Å, populated by `xftf`; `None` before calculation.
    pub fn get_r(&self) -> Option<ArrayBase<ViewRepr<&f64>, Ix1>> {
        Some(self.r.as_ref()?.view())
    }

    /// Read the stored value without recalculating.
    ///
    /// Full complex real-FFT coefficients, including DC and the even-length Nyquist
    /// bin. For dimensionless chi and k-weight w, units are Å⁻⁽ʷ⁺¹⁾. This cache is
    /// not truncated by `rmax_out`; `None` before calculation.
    pub fn get_chir(&self) -> Option<&DynRealDft<f64>> {
        self.chir.as_ref()
    }

    /// Return a newly allocated real component on the displayed R grid.
    /// Units are Å⁻⁽ʷ⁺¹⁾ for dimensionless chi. Returns `None` before calculation.
    pub fn get_chir_real(&self) -> Option<ArrayBase<OwnedRepr<f64>, Ix1>> {
        let len_r = self.r.as_ref()?.len();
        let chir = self.chir.as_ref()?;
        Some(Array1::from_iter(chir.iter().take(len_r).map(|x| x.re)))
    }

    /// Return a newly allocated imaginary component on the displayed R grid.
    /// Units are Å⁻⁽ʷ⁺¹⁾ for dimensionless chi. Returns `None` before calculation.
    pub fn get_chir_imag(&self) -> Option<ArrayBase<OwnedRepr<f64>, Ix1>> {
        let len_r = self.r.as_ref()?.len();
        let chir = self.chir.as_ref()?;
        Some(Array1::from_iter(chir.iter().take(len_r).map(|x| x.im)))
    }

    /// Read the stored value without recalculating.
    ///
    /// Magnitude of the displayed complex transform, paired with `r`.
    /// Units are Å⁻⁽ʷ⁺¹⁾ for dimensionless chi; `None` before calculation.
    pub fn get_chir_mag(&self) -> Option<ArrayBase<ViewRepr<&f64>, Ix1>> {
        Some(self.chir_mag.as_ref()?.view())
    }
    /// Read the stored value without recalculating.
    ///
    /// Dimensionless window on the prepared k grid, before k-weight multiplication.
    /// Use `Spectrum::kwin_k` to obtain the matching axis; `None` before calculation.
    pub fn get_kwin(&self) -> Option<ArrayBase<ViewRepr<&f64>, Ix1>> {
        Some(self.kwin.as_ref()?.view())
    }

    /// Read the stored value without recalculating.
    ///
    /// FFT sample spacing in Å⁻¹; `None` infers `k[1] - k[0]`. This controls both
    /// R spacing `pi / (nfft * kstep)` and amplitude `kstep / sqrt(pi)`.
    /// AUTOBK normally supplies 0.05 Å⁻¹; Input mode does not resample when this changes.
    pub fn get_kstep(&self) -> Option<&f64> {
        self.kstep.as_ref()
    }
}

/// Real inverse EXAFS transform with an R-space filter.
///
/// The inverse restores conjugate symmetry and applies `sqrt(pi)/(kstep*nfft)`
/// to an unnormalized inverse FFT. An all-pass filter with zero R-weight recovers
/// the forward weighted/windowed signal, not unweighted chi. This real-output
/// convention differs from Larch's complex analytic inverse.
///
/// For unchanged transform length N and sample spacing dk in Å⁻¹, let
/// `H_m=chi_R[m]*W_R(R_m)*R_m^u`, where u is the floored nonnegative R-weight.
/// The full N-bin sum includes the conjugate-symmetric negative frequencies:
///
/// ```text
/// chi_q[j] = sqrt(pi)/(N*dk) * sum(m=0..N-1, H_m * exp(2*pi*i*j*m/N))
/// q[j]     = j*dk
/// ```
///
/// q is in Å⁻¹ and chi_q has units Å^(u-w), where w was the forward k-weight.
/// With zero R-weight and an all-pass window, this recovers the forward
/// weighted/windowed input. A shell-selecting window discards information;
/// it does not undo absorption normalization or background removal.
/// Inferred kstep is stored. After changing nfft or the input R spacing, set
/// kstep=None or construct fresh settings before recalculating.
///
/// Change settings through `Spectrum::set_ifft` to invalidate dependent caches.
/// Direct field mutation requires an explicit `xftr` call.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(default)]
pub struct XrayFFTR {
    /// Maximum returned q in Å⁻¹; default 10.0, bounded by the inverse array length.
    pub qmax_out: Option<f64>,
    /// R-window family; default Kaiser–Bessel. `None` selects Hanning, not all-pass.
    pub window: Option<FTWindow>,
    /// Low-side R-window parameter; default 1.0. Usually a width in Å; Kaiser–Bessel
    /// also uses its numerical value for shape, and fractional Hanning uses fractions.
    pub dr: Option<f64>,
    /// High-side R-window parameter; `None` uses `dr`. Its meaning depends on the family.
    pub dr2: Option<f64>,
    /// Lower R-window boundary in Å; default 0.0. `None` uses the first supplied R.
    pub rmin: Option<f64>,
    /// Upper R-window boundary in Å; default 20.0. `None` uses the last supplied R.
    pub rmax: Option<f64>,
    /// Additional power of R applied before inversion; default 0.0.
    /// Nonnegative fractional values are floored. This does not remove forward k-weighting.
    pub rweight: Option<f64>,
    /// Inverse FFT length; default 2048. Changing it resizes the positive-frequency
    /// coefficients and changes automatic q spacing; leave `kstep` automatic.
    pub nfft: Option<usize>,
    /// Output q spacing in Å⁻¹. `None` infers `pi / (R_step * nfft)`.
    /// An explicit positive value must agree with the supplied uniform R grid.
    pub kstep: Option<f64>,
    /// Returned q grid in Å⁻¹, beginning at zero; `None` before calculation.
    pub q: Option<ArrayBase<OwnedRepr<f64>, Ix1>>,
    /// Full real inverse signal before the `qmax_out` display limit. It retains
    /// forward weighting/windowing. Units are Å⁽ᵘ⁻ʷ⁾ for dimensionless input chi,
    /// forward k-weight w and inverse R-weight u; `None` before calculation.
    pub chiq: Option<ArrayBase<OwnedRepr<f64>, Ix1>>,
    /// Full R-window multiplied by `R.powf(rweight)`, paired with all positive
    /// forward bins, not just the displayed R range. It is dimensionless only when
    /// `rweight=0`; `None` before calculation.
    pub rwin: Option<ArrayBase<OwnedRepr<f64>, Ix1>>,
}

impl Default for XrayFFTR {
    fn default() -> Self {
        XrayFFTR {
            qmax_out: Some(10.0),
            window: Some(FTWindow::KaiserBessel),
            dr: Some(1.),
            dr2: None,
            rmin: Some(0.),
            rmax: Some(20.0),
            rweight: Some(0.0),
            nfft: Some(2048),
            kstep: None,
            q: None,
            chiq: None,
            rwin: None,
        }
    }
}

impl XrayFFTR {
    /// Create recommended inverse settings with automatic q spacing.
    pub fn new() -> XrayFFTR {
        XrayFFTR::default()
    }

    /// Resolve automatic settings from a nonempty r grid without calculating.
    ///
    /// Prefer the validated transform method for public input. This lower-level
    /// helper indexes the supplied grid directly and can panic on an empty grid.
    /// Only None fields are resolved; their resulting values remain stored when
    /// later inputs change. Reset them to None to infer from a new grid.
    pub fn fill_parameter(&mut self, r: ArrayBase<ViewRepr<&f64>, Ix1>) -> &mut Self {
        if self.rweight.is_none() {
            self.rweight = Some(0.0);
        }

        self.rweight = Some(self.rweight.unwrap().max(0.0).floor());

        if self.rmin.is_none() {
            self.rmin = Some(r[0]);
        }

        if self.rmax.is_none() {
            self.rmax = Some(r[r.len() - 1]);
        }

        if self.dr.is_none() {
            self.dr = Some(1.0);
        }

        if self.nfft.is_none() {
            self.nfft = Some(2048);
        }

        if self.qmax_out.is_none() {
            self.qmax_out = Some(10.0);
        }

        if self.kstep.is_none() {
            self.kstep = Some(std::f64::consts::PI / (r[1] - r[0]) / self.nfft.unwrap() as f64);
        }
        self
    }

    /// Validate inverse settings and allocate `(filtered_coefficients, R_filter)`.
    /// The filter includes `R.powf(rweight)` and all positive-frequency bins,
    /// independently of the displayed R length. Only unset settings are
    /// resolved; no inverse signal is calculated or cached by this step.
    pub fn xftr_prep(
        &mut self,
        r: ArrayBase<ViewRepr<&f64>, Ix1>,
        chir: &DynRealDft<f64>,
    ) -> Result<(DynRealDft<f64>, ArrayBase<OwnedRepr<f64>, Ix1>), FFTError> {
        super::inverse_fft::validate(&r.to_vec(), self)?;
        self.fill_parameter(r);
        let rstep = std::f64::consts::PI / self.kstep.unwrap() / self.nfft.unwrap() as f64;
        let full_r = Array1::from_iter((0..chir.len()).map(|i| i as f64 * rstep));
        let mut win = ftwindow(
            &full_r,
            self.rmin,
            self.rmax,
            self.dr,
            self.dr2,
            self.window,
        )
        .map_err(|e| FFTError::WindowCalculationFailed {
            reason: e.to_string(),
        })?;
        let weight = self.rweight.unwrap();
        if weight > 0.0 {
            win *= &full_r.mapv(|radius| radius.powf(weight));
        }
        let chir_win = chir * win.as_slice().expect("owned contiguous inverse window");

        Ok((chir_win, win))
    }

    /// Filter full complex forward coefficients and cache a real inverse signal.
    ///
    /// `r` must contain at least two uniformly increasing finite samples starting
    /// at zero. Explicit `kstep` and `nfft` must match its spacing. Invalid grids,
    /// negative/nonfinite parameters or `rmin >= rmax` return `FFTError`. The input
    /// coefficients are borrowed unchanged; `rmax_out` never truncates this filter.
    pub fn xftr(
        &mut self,
        r: ArrayBase<ViewRepr<&f64>, Ix1>,
        chir: &DynRealDft<f64>,
    ) -> Result<&mut Self, FFTError> {
        let (chir_win, win) = self.xftr_prep(r, chir)?;
        let nfft = self.nfft.unwrap();
        let out = xftr_fast(&chir_win, nfft, self.kstep.unwrap());

        let q = Array1::from_vec(super::inverse_fft::q_grid(
            self.qmax_out.unwrap(),
            self.kstep.unwrap(),
            out.len(),
        ));

        self.q = Some(q);
        self.rwin = Some(win);
        self.chiq = Some(out);

        Ok(self)
    }

    /// Read the stored value without recalculating.
    ///
    /// Returned q grid in Å⁻¹, beginning at zero; `None` before calculation.
    pub fn get_q(&self) -> Option<ArrayBase<ViewRepr<&f64>, Ix1>> {
        Some(self.q.as_ref()?.view())
    }

    /// Return a newly allocated real inverse signal limited to the returned q grid.
    /// Forward k-weighting and windowing remain; returns `None` before calculation.
    pub fn get_chiq(&self) -> Option<ArrayBase<OwnedRepr<f64>, Ix1>> {
        let len_q = self.q.as_ref()?.len();

        Some(
            self.chiq
                .clone()?
                .slice_axis(Axis(0), (0..len_q).into())
                .to_owned(),
        )
    }

    /// Read the stored value without recalculating.
    ///
    /// Full R-window multiplied by `R.powf(rweight)`, paired with all positive
    /// forward bins, not just the displayed R range. It is dimensionless only when
    /// `rweight=0`; `None` before calculation.
    pub fn get_rwin(&self) -> Option<ArrayBase<ViewRepr<&f64>, Ix1>> {
        Some(self.rwin.as_ref()?.view())
    }

    /// Read the stored value without recalculating.
    ///
    /// Output q spacing in Å⁻¹. `None` infers `pi / (R_step * nfft)`.
    /// An explicit positive value must agree with the supplied uniform R grid.
    pub fn get_kstep(&self) -> Option<&f64> {
        self.kstep.as_ref()
    }

    /// Read the stored value without recalculating.
    ///
    /// Additional power of R applied before inversion; default 0.0.
    /// Nonnegative fractional values are floored. This does not remove forward k-weighting.
    pub fn get_rweight(&self) -> Option<&f64> {
        self.rweight.as_ref()
    }

    /// Read the stored value without recalculating.
    ///
    /// Inverse FFT length; default 2048. Changing it resizes the positive-frequency
    /// coefficients and changes automatic q spacing; leave `kstep` automatic.
    pub fn get_nfft(&self) -> Option<&usize> {
        self.nfft.as_ref()
    }

    /// Read the stored value without recalculating.
    ///
    /// R-window family; default Kaiser–Bessel. `None` selects Hanning, not all-pass.
    pub fn get_window(&self) -> Option<&FTWindow> {
        self.window.as_ref()
    }
}

/// Transform already weighted/windowed ndarray input using `kstep / sqrt(pi)`.
/// Uses an unnormalized negative-exponent FFT and zero-pads or truncates to
/// `nfft`. This low-level helper does not validate parameters.
pub fn xftf_fast(chi: ArrayBase<ViewRepr<&f64>, Ix1>, nfft: usize, kstep: f64) -> DynRealDft<f64> {
    let mut cchi = vec![0.0_f64; nfft];
    for (output, input) in cchi.iter_mut().zip(chi.iter()) {
        *output = *input;
    }

    let mut freq = cchi.real_fft();

    freq *= kstep / (std::f64::consts::PI).sqrt();

    freq
}

/// Restore conjugate symmetry and return a real ndarray inverse signal.
/// The unnormalized inverse is multiplied by `sqrt(pi)/(kstep*nfft)`;
/// this low-level helper does not validate parameters.
pub fn xftr_fast(
    chir: &DynRealDft<f64>,
    nfft: usize,
    kstep: f64,
) -> ArrayBase<OwnedRepr<f64>, Ix1> {
    Array1::from_vec(super::inverse_fft::inverse(chir, nfft, kstep))
}

/// Transform an already weighted/windowed real signal with `kstep / sqrt(pi)`.
///
/// Uses an unnormalized negative-exponent FFT and returns all positive-frequency
/// bins. The input is zero-padded or truncated to `nfft`. This low-level helper
/// does not validate sizes or spacing; prefer `XrayFFTF::xftf` for checked input.
pub fn xftf_fast_nalgebra(chi: &DVector<f64>, nfft: usize, kstep: f64) -> DynRealDft<f64> {
    let mut cchi = vec![0.0_f64; nfft];
    let length = chi.len().min(nfft);
    cchi[..length].copy_from_slice(&chi.as_slice()[..length]);

    let mut freq = cchi.real_fft();

    freq *= kstep / std::f64::consts::PI.sqrt();

    freq
}

/// Invert already filtered positive-frequency coefficients into a real signal.
///
/// Restores conjugate symmetry, resizes to `nfft`, and multiplies the unnormalized
/// inverse by `sqrt(pi)/(kstep*nfft)`. This helper does not validate settings; use
/// `XrayFFTR::xftr` for checked input.
pub fn xftr_fast_nalgebra(chir: &DynRealDft<f64>, nfft: usize, kstep: f64) -> DVector<f64> {
    DVector::from_vec(super::inverse_fft::inverse(chir, nfft, kstep))
}

/// Convenience extension for the unnormalized forward FFT with XAFS scaling.
pub trait XFFT {
    /// Transform already weighted/windowed input; see the corresponding fast helper.
    fn xftf_fast(&self, nfft: usize, kstep: f64) -> DynRealDft<f64>;
}

impl XFFT for ArrayBase<OwnedRepr<f64>, Ix1> {
    fn xftf_fast(&self, nfft: usize, kstep: f64) -> DynRealDft<f64> {
        xftf_fast(self.view(), nfft, kstep)
    }
}

impl XFFT for ArrayBase<ViewRepr<&f64>, Ix1> {
    fn xftf_fast(&self, nfft: usize, kstep: f64) -> DynRealDft<f64> {
        xftf_fast(*self, nfft, kstep)
    }
}

impl XFFT for DVector<f64> {
    fn xftf_fast(&self, nfft: usize, kstep: f64) -> DynRealDft<f64> {
        xftf_fast_nalgebra(self, nfft, kstep)
    }
}

/// Convenience extension for a real inverse transform with XAFS scaling.
pub trait XFFTReverse<T> {
    /// Invert already filtered positive-frequency coefficients into a real array.
    fn xftr_fast(&self, nfft: usize, kstep: f64) -> T;
}

impl XFFTReverse<ArrayBase<OwnedRepr<f64>, Ix1>> for DynRealDft<f64> {
    fn xftr_fast(&self, nfft: usize, kstep: f64) -> ArrayBase<OwnedRepr<f64>, Ix1> {
        xftr_fast(self, nfft, kstep)
    }
}

impl XFFTReverse<DVector<f64>> for DynRealDft<f64> {
    fn xftr_fast(&self, nfft: usize, kstep: f64) -> DVector<f64> {
        xftr_fast_nalgebra(self, nfft, kstep)
    }
}

/// Allocate component arrays from complex Fourier coefficients.
pub trait FFTUtils<T> {
    /// Allocate interleaved real and imaginary values `[re0, im0, re1, im1, ...]`.
    /// Allocate the real component of every coefficient.
    fn realimg(&self) -> T;
    /// Allocate the real component of every coefficient.
    fn re(&self) -> T;
    /// Allocate the imaginary component of every coefficient.
    fn im(&self) -> T;
    /// Allocate each coefficient magnitude `sqrt(re*re + im*im)`.
    fn norm(&self) -> T;
    /// Allocate each coefficient magnitude `sqrt(re*re + im*im)`.
    /// Allocate each squared magnitude `re*re + im*im`.
    fn norm_sqr(&self) -> T;
}

impl FFTUtils<ArrayBase<OwnedRepr<f64>, Ix1>> for DynRealDft<f64> {
    fn realimg(&self) -> Array1<f64> {
        Array1::from_iter(self.iter().flat_map(|x| vec![x.re, x.im]))
    }

    fn re(&self) -> Array1<f64> {
        Array1::from_iter(self.iter().map(|x| x.re))
    }

    fn im(&self) -> Array1<f64> {
        Array1::from_iter(self.iter().map(|x| x.im))
    }

    fn norm(&self) -> Array1<f64> {
        Array1::from_iter(self.iter().map(|x| x.norm()))
    }

    fn norm_sqr(&self) -> Array1<f64> {
        Array1::from_iter(self.iter().map(|x| x.norm_sqr()))
    }
}

impl FFTUtils<DVector<f64>> for DynRealDft<f64> {
    fn realimg(&self) -> DVector<f64> {
        DVector::from_iterator(self.len() * 2, self.iter().flat_map(|x| vec![x.re, x.im]))
    }

    fn re(&self) -> DVector<f64> {
        DVector::from_iterator(self.len(), self.iter().map(|x| x.re))
    }

    fn im(&self) -> DVector<f64> {
        DVector::from_iterator(self.len(), self.iter().map(|x| x.im))
    }

    fn norm(&self) -> DVector<f64> {
        DVector::from_iterator(self.len(), self.iter().map(|x| x.norm()))
    }

    fn norm_sqr(&self) -> DVector<f64> {
        DVector::from_iterator(self.len(), self.iter().map(|x| x.norm_sqr()))
    }
}

impl FFTUtils<ArrayBase<OwnedRepr<f64>, Ix1>> for [Complex<f64>] {
    fn realimg(&self) -> Array1<f64> {
        Array1::from_iter(self.iter().flat_map(|x| vec![x.re, x.im]))
    }

    fn re(&self) -> Array1<f64> {
        Array1::from_iter(self.iter().map(|x| x.re))
    }

    fn im(&self) -> Array1<f64> {
        Array1::from_iter(self.iter().map(|x| x.im))
    }

    fn norm(&self) -> Array1<f64> {
        Array1::from_iter(self.iter().map(|x| x.norm()))
    }

    fn norm_sqr(&self) -> Array1<f64> {
        Array1::from_iter(self.iter().map(|x| x.norm_sqr()))
    }
}

impl FFTUtils<DVector<f64>> for [Complex<f64>] {
    fn realimg(&self) -> DVector<f64> {
        DVector::from_iterator(self.len() * 2, self.iter().flat_map(|x| vec![x.re, x.im]))
    }

    fn re(&self) -> DVector<f64> {
        DVector::from_iterator(self.len(), self.iter().map(|x| x.re))
    }

    fn im(&self) -> DVector<f64> {
        DVector::from_iterator(self.len(), self.iter().map(|x| x.im))
    }

    fn norm(&self) -> DVector<f64> {
        DVector::from_iterator(self.len(), self.iter().map(|x| x.norm()))
    }

    fn norm_sqr(&self) -> DVector<f64> {
        DVector::from_iterator(self.len(), self.iter().map(|x| x.norm_sqr()))
    }
}

// #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
// pub struct DynRealDftWrapper {
//     zeroth_bin: f64,

//     frequency_bins: Vec<Complex<f64>>,

//     original_length: usize,
// }

// impl From<DynRealDft<f64>> for DynRealDftWrapper {
//     fn from(dft: DynRealDft<f64>) -> Self {
//         DynRealDftWrapper {
//             zeroth_bin: dft.get_offset().clone(),
//             frequency_bins: dft.get_frequency_bins().to_vec(),
//             original_length: dft.get_original_length(),
//         }
//     }
// }

// impl PartialEq for DynRealDft<f64> {
//     fn eq(&self, other: &Self) -> bool {
//         self.len() == other.len()
//             && self.get_offset() == other.get_offset()
//             && self.get_frequency_bins() == other.get_frequency_bins()
//     }
// }

// #[derive(Debug, Clone, PartialEq, Default)]
// pub struct XrayFFTR {}

#[cfg(test)]
mod tests {
    use easyfft::prelude::*;
    use ndarray::Array1;

    use super::*;
    use crate::xafs::io;
    use crate::xafs::nshare::ToNalgebra;
    use approx::{assert_abs_diff_eq, assert_relative_eq};

    use data_reader::reader::{load_txt_f64, Delimiter, ReaderParams};

    use crate::xafs::tests::PARAM_LOADTXT;
    use crate::xafs::tests::TEST_TOL;
    use crate::xafs::tests::TOP_DIR;

    use crate::xafs::background::BackgroundMethod;
    use crate::xafs::background::AUTOBK;

    const ACCEPTABLE_MU_DIFF: f64 = 1e-6;
    const CHI_MSE_TOL: f64 = 1e-2;
    const CHI_Q_TOL: f64 = 1e-1;

    #[test]
    fn test_xftf_fast() {
        let x: Array1<f64> = Array1::linspace(0., 10., 10);
        let sin_x = x.map(|x| x.sin());
        let nfft = 10;
        let kstep = 1.;
        let fft = xftf_fast(sin_x.view(), nfft, kstep);

        let norm: DVector<f64> = fft.norm();

        let expected_norm = DVector::from(vec![
            0.6822515304148188,
            0.999816632055004,
            2.4133321684349966,
            0.35447122637608214,
            0.16620199767343982,
            0.1252841340812192,
        ]);

        assert_abs_diff_eq!(norm, expected_norm, epsilon = 1e-16);
    }

    #[test]
    fn test_xftr_fast() {
        let x: Array1<f64> = Array1::linspace(0., 10., 1024);
        let sin_x = x.map(|x| x.sin());
        let nfft = 1024;
        let kstep = 0.1;

        let fft = xftf_fast(sin_x.view(), nfft, kstep);
        let ifft = xftr_fast(&fft, nfft, kstep);

        sin_x.iter().zip(ifft.iter()).for_each(|(x, y)| {
            assert_abs_diff_eq!(x, y, epsilon = 1e-12);
        });
    }

    #[test]
    fn test_xftf_fast_nalgebra() {
        let x: DVector<f64> = Array1::linspace(0., 10., 10).into_nalgebra();
        let sin_x = x.map(|x| x.sin());
        let nfft = 10;
        let kstep = 1.;
        let fft = xftf_fast_nalgebra(&sin_x, nfft, kstep);

        let norm: DVector<f64> = fft.norm();

        let expected_norm = DVector::from(vec![
            0.6822515304148188,
            0.999816632055004,
            2.4133321684349966,
            0.35447122637608214,
            0.16620199767343982,
            0.1252841340812192,
        ]);

        assert_abs_diff_eq!(norm, expected_norm, epsilon = 1e-16);
    }

    #[test]
    fn test_xftr_fast_nalgebra() {
        let x: DVector<f64> = Array1::linspace(0., 10., 1024).into_nalgebra();
        let sin_x = x.map(|x| x.sin());
        let nfft = 1024;
        let kstep = 0.1;

        let fft = xftf_fast_nalgebra(&sin_x, nfft, kstep);
        let ifft = xftr_fast_nalgebra(&fft, nfft, kstep);

        sin_x.iter().zip(ifft.iter()).for_each(|(x, y)| {
            assert_abs_diff_eq!(x, y, epsilon = 1e-12);
        });
    }

    #[test]
    #[allow(non_snake_case)]
    fn test_Xray_FFTF() -> Result<(), Box<dyn std::error::Error>> {
        let path = String::from(TOP_DIR) + "/tests/testfiles/Ru_QAS.dat";
        let mut xafs_test_group = io::load_spectrum_QAS_trans(&path).unwrap();

        xafs_test_group.set_background_method(Some(BackgroundMethod::AUTOBK(AUTOBK {
            rbkg: Some(1.4),
            kweight: Some(2),
            ..Default::default()
        })))?;
        xafs_test_group.calc_background()?;

        xafs_test_group.xftf = Some(XrayFFTF {
            window: Some(FTWindow::Hanning),
            dk: Some(1.0),
            kmin: Some(2.0),
            kmax: Some(15.0),
            kweight: Some(2.0),
            ..Default::default()
        });
        xafs_test_group.fft()?;

        let larch_r_path = String::from(TOP_DIR) + "/tests/testfiles/Ru_QAS_xftf_larch.txt";
        let larch_r = load_txt_f64(&larch_r_path, &PARAM_LOADTXT).unwrap();

        let r_expected = larch_r.get_col(0);
        let chir_expected = larch_r.get_col(1);

        let r = xafs_test_group.r().unwrap();
        let chir = xafs_test_group.chir_mag().unwrap();

        r.iter().zip(r_expected.iter()).for_each(|(x, y)| {
            assert_abs_diff_eq!(x, y, epsilon = TEST_TOL);
        });

        let mse = chir
            .iter()
            .zip(chir_expected.iter())
            .map(|(x, y)| (x - y).powi(2))
            .sum::<f64>()
            / r.len() as f64;

        // println!("mse: {}", mse);
        assert!(mse < CHI_MSE_TOL);

        // chir.iter().zip(chir_expected.iter()).for_each(|(x, y)| {
        //     println!("[{}, {}],", x, y);
        // });

        // mse = r
        // .iter()
        // .zip(r_expected.iter())
        // .map(|(x, y)| (x - y).powi(2))
        // .sum::<f64>()
        // / r.len() as f64;

        // println!("mse: {}", mse);

        // assert!(mse < CHI_MSE_TOL);

        Ok(())
    }

    #[test]
    fn test_chir_real_imag_are_distinct() -> Result<(), Box<dyn std::error::Error>> {
        let path = String::from(TOP_DIR) + "/tests/testfiles/Ru_QAS.dat";
        let mut xafs_test_group = io::load_spectrum_QAS_trans(&path).unwrap();

        xafs_test_group.set_background_method(Some(BackgroundMethod::AUTOBK(AUTOBK {
            rbkg: Some(1.4),
            kweight: Some(2),
            ..Default::default()
        })))?;
        xafs_test_group.calc_background()?;

        xafs_test_group.xftf = Some(XrayFFTF {
            window: Some(FTWindow::Hanning),
            dk: Some(1.0),
            kmin: Some(2.0),
            kmax: Some(15.0),
            kweight: Some(2.0),
            ..Default::default()
        });
        xafs_test_group.fft()?;

        let chir_real = xafs_test_group.chir_real().unwrap();
        let chir_imag = xafs_test_group.chir_imag().unwrap();
        assert_eq!(chir_real.len(), chir_imag.len());

        let has_distinct_value = chir_real
            .iter()
            .zip(chir_imag.iter())
            .any(|(re, im)| (re - im).abs() > TEST_TOL);
        assert!(has_distinct_value);

        Ok(())
    }

    #[test]
    #[allow(non_snake_case)]
    fn test_XrayFFTR() -> Result<(), Box<dyn std::error::Error>> {
        let path = String::from(TOP_DIR) + "/tests/testfiles/Ru_QAS.dat";
        let mut xafs_test_group = io::load_spectrum_QAS_trans(&path).unwrap();

        xafs_test_group.set_background_method(Some(BackgroundMethod::AUTOBK(AUTOBK {
            rbkg: Some(1.4),
            ..Default::default()
        })))?;
        xafs_test_group.calc_background()?;

        xafs_test_group.xftf = Some(XrayFFTF {
            window: Some(FTWindow::Hanning),
            dk: Some(f64::EPSILON),
            kmin: Some(0.0),
            kmax: Some(15.0),
            kweight: Some(2.0),
            ..Default::default()
        });
        xafs_test_group.fft()?;

        xafs_test_group.xftr = Some(XrayFFTR {
            window: Some(FTWindow::Hanning),
            rweight: Some(0.0),
            dr: Some(f64::EPSILON),
            rmin: Some(0.0),
            rmax: Some(10.0),
            ..Default::default()
        });
        xafs_test_group.ifft()?;

        let q = xafs_test_group.q().unwrap();
        let chiq = xafs_test_group.chiq().unwrap();

        // println!("q: {:?}", q.len());
        // println!("chiq: {:?}", chiq.len());

        assert!(q.len() == chiq.len());

        // Convert DVector to Array1 for multiplication with Array1 view
        let chi_dv = xafs_test_group.chi_kweighted().unwrap();
        let chi_array = Array1::from_vec(chi_dv.data.as_vec().clone());
        let chi = chi_array * xafs_test_group.xftf.unwrap().get_kwin().unwrap();

        let chi = chi.slice_axis(Axis(0), (0..chi.len()).into()).to_owned();

        chi.iter().zip(chiq.iter()).for_each(|(x, y)| {
            assert_relative_eq!(x, y, epsilon = CHI_Q_TOL);
        });

        Ok(())
    }
}