spectrograms 2.1.3

//! 2D FFT operations for image and spatial data processing.
//!
//! This module provides 2D Fast Fourier Transform operations optimized for
//! image processing, spatial filtering, and 2D signal analysis.
//!
//! # Overview
//!
//! The 2D FFT transforms a 2D real-valued array (like an image) into the
//! frequency domain, enabling efficient convolution, filtering, and analysis.
//!
//! # Examples
//!
//! ```rust
//! use spectrograms::fft2d::fft2d;
//! use ndarray::Array2;
//!
//! // Create a 128x128 image
//! let image = Array2::<f64>::zeros((128, 128));
//!
//! // Compute 2D FFT
//! let spectrum = fft2d(&image.view()).unwrap();
//!
//! // Output shape is 128 x 65 due to Hermitian symmetry
//! assert_eq!(spectrum.shape(), &[128, 65]);
//! ```
//!
//! # Output Shape
//!
//! Due to Hermitian symmetry in real-to-complex FFTs, the output shape is
//! `(nrows, ncols/2 + 1)` for an input of shape `(nrows, ncols)`.

use std::collections::HashMap;
use std::collections::hash_map::Entry;

use ndarray::{Array2, ArrayView2};
use num_complex::Complex;

use crate::Sample;
use crate::fft_backend::{C2rPlan2d, R2cPlan2d, r2c_output_size_2d};
use crate::{SpectrogramError, SpectrogramResult};

/// Compute 2D FFT of a real-valued array (e.g., image).
///
/// Transforms a 2D real array into the frequency domain using a real-to-complex
/// FFT. The output exploits Hermitian symmetry, so only half the frequencies
/// are stored along the column dimension.
///
/// # Arguments
///
/// * `data` - Input 2D array with shape `(nrows, ncols)`
///
/// # Returns
///
/// Complex 2D array with shape `(nrows, ncols/2 + 1)` containing the frequency
/// domain representation.
///
/// # Errors
///
/// Returns `SpectrogramError` if the input array is invalid or processing fails.
///
/// # Examples
///
/// ```rust
/// use spectrograms::fft2d::fft2d;
/// use ndarray::Array2;
///
/// let img = Array2::<f64>::zeros((64, 64));
/// let spectrum = fft2d(&img.view()).unwrap();
/// assert_eq!(spectrum.shape(), &[64, 33]); // 64 x (64/2 + 1)
/// ```
///
/// # Performance
///
/// For batch processing of multiple arrays with the same dimensions, use
/// [`Fft2dPlanner`] to reuse FFT plans and avoid repeated setup overhead.
#[inline]
pub fn fft2d<T: Sample>(data: &ArrayView2<T>) -> SpectrogramResult<Array2<Complex<T>>> {
    let (nrows, ncols) = (data.nrows(), data.ncols());

    if nrows == 0 || ncols == 0 {
        return Err(SpectrogramError::invalid_input(
            "array dimensions must be > 0",
        ));
    }

    // Ensure array is contiguous and row-major
    // Tricky to convert it here due to ownership and cloning - better to require caller to provide standard layout
    if !data.is_standard_layout() {
        return Err(SpectrogramError::invalid_input(
            "array must be contiguous and row-major (standard layout)",
        ));
    }

    let out_shape = r2c_output_size_2d(nrows, ncols);

    let mut plan = T::plan_r2c_2d(nrows, ncols)?;

    let input_slice = data
        .as_slice()
        .ok_or_else(|| SpectrogramError::invalid_input("array must be contiguous"))?;

    let mut output = vec![Complex::new(T::zero(), T::zero()); out_shape.0 * out_shape.1];

    plan.process(input_slice, &mut output)?;

    Array2::from_shape_vec(out_shape, output)
        .map_err(|e| SpectrogramError::invalid_input(format!("failed to reshape output: {e}")))
}

/// Compute inverse 2D FFT from frequency domain back to spatial domain.
///
/// Transforms a complex 2D frequency array back into a real spatial array.
/// This is the inverse operation of [`fft2d`].
///
/// # Arguments
///
/// * `spectrum` - Complex frequency array with shape `(nrows, ncols/2 + 1)`
/// * `output_ncols` - Number of columns in the output (must be even or odd as per original)
///
/// # Returns
///
/// Real 2D array with shape `(nrows, output_ncols)`.
///
/// # Errors
///
/// Returns `SpectrogramError` if the input array is invalid or processing fails.
///
/// # Examples
///
/// ```rust
/// use spectrograms::fft2d::{fft2d, ifft2d};
/// use spectrograms::nzu;
/// use ndarray::Array2;
///
/// let original = Array2::<f64>::from_shape_fn((64, 64), |(i, j)| {
///     (i as f64 * 0.1).sin() + (j as f64 * 0.2).cos()
/// });
///
/// let spectrum = fft2d(&original.view()).unwrap();
/// let reconstructed = ifft2d(&spectrum, 64).unwrap();
///
/// // Check roundtrip accuracy
/// for ((i, j), &val) in original.indexed_iter() {
///     assert!((reconstructed[[i, j]] - val).abs() < 1e-10);
/// }
/// ```
#[inline]
pub fn ifft2d<T: Sample>(
    spectrum: &Array2<Complex<T>>,
    output_ncols: usize,
) -> SpectrogramResult<Array2<T>> {
    let nrows = spectrum.nrows();

    if nrows == 0 || output_ncols == 0 {
        return Err(SpectrogramError::invalid_input("dimensions must be > 0"));
    }

    // Verify input shape matches expected frequency domain shape
    let expected_ncols = output_ncols / 2 + 1;
    if spectrum.ncols() != expected_ncols {
        return Err(SpectrogramError::dimension_mismatch(
            expected_ncols,
            spectrum.ncols(),
        ));
    }

    if !spectrum.is_standard_layout() {
        return Err(SpectrogramError::invalid_input(
            "array must be contiguous and row-major (standard layout)",
        ));
    }

    let mut plan = T::plan_c2r_2d(nrows, output_ncols)?;

    let input_slice = spectrum
        .as_slice()
        .ok_or_else(|| SpectrogramError::invalid_input("array must be contiguous"))?;

    let mut output = vec![T::zero(); nrows * output_ncols];

    plan.process(input_slice, &mut output)?;

    Array2::from_shape_vec((nrows, output_ncols), output)
        .map_err(|e| SpectrogramError::invalid_input(format!("failed to reshape output: {e}")))
}

/// Compute 2D power spectrum (squared magnitude).
///
/// # Arguments
///
/// * `data` - Input 2D array with shape `(nrows, ncols)`
///
/// # Returns
/// Returns `|FFT(x)|²` for each frequency component, useful for analyzing
/// energy distribution across frequencies.
///
/// # Errors
///
/// Returns `SpectrogramError` if the input array is invalid or processing fails.
///
/// # Examples
///
/// ```rust
/// use spectrograms::fft2d::power_spectrum_2d;
/// use ndarray::Array2;
///
/// let img = Array2::<f64>::ones((32, 32));
/// let power = power_spectrum_2d(&img.view()).unwrap();
///
/// // DC component (0,0) should have all the energy
/// assert!(power[[0, 0]] > 1000.0);
/// ```
#[inline]
pub fn power_spectrum_2d<T: Sample>(data: &ArrayView2<T>) -> SpectrogramResult<Array2<T>> {
    let spectrum = fft2d(data)?;
    let power = spectrum.mapv(|c| c.norm_sqr());
    Ok(power)
}

/// Compute 2D magnitude spectrum.
///
/// # Arguments
///
/// * `data` - Input 2D array with shape `(nrows, ncols)`
///
/// # Returns
///
/// Real 2D array with shape `(nrows, ncols/2 + 1)` containing the magnitude spectrum.
///
/// # Errors
///
/// Returns `SpectrogramError` if the input array is invalid or processing fails.
///
/// # Examples
///
/// ```rust
/// use spectrograms::fft2d::magnitude_spectrum_2d;
/// use ndarray::Array2;
///
/// let img = Array2::<f64>::zeros((32, 32));
/// let magnitude = magnitude_spectrum_2d(&img.view()).unwrap();
/// ```
#[inline]
pub fn magnitude_spectrum_2d<T: Sample>(data: &ArrayView2<T>) -> SpectrogramResult<Array2<T>> {
    let spectrum = fft2d(data)?;
    let magnitude = spectrum.mapv(num_complex::Complex::norm);
    Ok(magnitude)
}

/// Shift zero-frequency component to center (like `numpy.fft.fftshift`).
///
/// This rearranges the FFT output so that the DC (zero-frequency) component
/// is at the center of the array, which is useful for visualization.
///
/// # Arguments
///
/// * `arr` - Input 2D array to shift
///
/// # Returns
///
/// 2D array with quadrants swapped to center the zero-frequency component.
///
/// # Note
///
/// This function works on the **full** FFT output. Since [`fft2d`] returns
/// a reduced spectrum due to Hermitian symmetry, you may need to expand it
/// first for a full centered view, or use this on magnitude/power spectra.
///
/// # Examples
///
/// ```rust
/// use spectrograms::fft2d::{power_spectrum_2d, fftshift};
/// use ndarray::{Array2, ArrayView2};
///
/// let img = Array2::<f64>::ones((32, 32));
/// let power = power_spectrum_2d(&img.view()).unwrap();
/// let shifted = fftshift(power);
/// ```
#[inline]
#[must_use]
pub fn fftshift<T: Clone>(arr: Array2<T>) -> Array2<T> {
    let (nrows, ncols) = arr.dim();
    if nrows == 0 || ncols == 0 {
        return arr;
    }

    let row_shift = nrows / 2;
    let col_shift = ncols / 2;

    shift_2d(arr, row_shift, col_shift)
}

/// Inverse of [`fftshift`] - shift center back to corners.
///
/// Undoes the frequency shift performed by [`fftshift`].
///
/// # Arguments
///
/// * `arr` - Input 2D array to inverse shift
///
/// # Returns
///
/// 2D array with quadrants swapped back to original positions.
#[inline]
#[must_use]
pub fn ifftshift<T: Clone>(arr: Array2<T>) -> Array2<T> {
    let (nrows, ncols) = arr.dim();
    if nrows == 0 || ncols == 0 {
        return arr;
    }

    let row_shift = nrows.div_ceil(2);
    let col_shift = ncols.div_ceil(2);

    shift_2d(arr, row_shift, col_shift)
}

fn shift_2d<T: Clone>(arr: Array2<T>, row_shift: usize, col_shift: usize) -> Array2<T> {
    let (nrows, ncols) = arr.dim();
    if nrows == 0 || ncols == 0 {
        return arr;
    }

    let row_shift = row_shift % nrows;
    let col_shift = col_shift % ncols;

    if row_shift == 0 && col_shift == 0 {
        return arr;
    }

    Array2::from_shape_fn((nrows, ncols), |(i, j)| {
        let src_i = (i + row_shift) % nrows;
        let src_j = (j + col_shift) % ncols;
        arr[[src_i, src_j]].clone()
    })
}

/// Shift zero-frequency component to center for 1D array.
///
/// # Arguments
///
/// * `arr` - Input 1D array to shift
///
/// # Returns
///
/// 1D array with zero-frequency component at center.
#[inline]
#[must_use]
pub fn fftshift_1d<T: Clone>(arr: Vec<T>) -> Vec<T> {
    let len = arr.len();
    rotate_left(arr, len / 2)
}

/// Inverse of fftshift for 1D arrays.
///
/// # Arguments
///
/// * `arr` - Input 1D array to inverse shift
///
/// # Returns
///
/// 1D array with quadrants swapped back to original positions.
#[inline]
#[must_use]
pub fn ifftshift_1d<T: Clone>(arr: Vec<T>) -> Vec<T> {
    let len = arr.len();
    rotate_left(arr, len.div_ceil(2))
}

fn rotate_left<T: Clone>(arr: Vec<T>, shift: usize) -> Vec<T> {
    let n = arr.len();
    if n == 0 {
        return arr;
    }
    let k = shift % n;
    if k == 0 {
        return arr;
    }

    let mut result = Vec::with_capacity(n);
    result.extend_from_slice(&arr[k..]);
    result.extend_from_slice(&arr[..k]);
    result
}

/// Compute FFT sample frequencies (like `numpy.fft.fftfreq`).
///
/// Returns the sample frequencies (in cycles per unit of the sample spacing) for FFT output.
/// For an FFT of length `n` with sample spacing `d`:
/// - Frequencies go from 0 to (n-1)//(2*n*d) for positive frequencies
/// - Then wrap to negative frequencies
///
/// # Arguments
///
/// * `n` - Window length (number of samples)
/// * `d` - Sample spacing (inverse of sampling rate). Default is 1.0.
///
/// # Returns
///
/// Vector of length `n` containing the frequency bin centers in cycles per unit.
///
/// # Examples
///
/// ```rust
/// use spectrograms::fft2d::fftfreq;
///
/// // For 8 samples with spacing 1.0
/// let freqs = fftfreq::<f64>(8, 1.0);
/// // Returns: [0.0, 0.125, 0.25, 0.375, -0.5, -0.375, -0.25, -0.125]
///
/// // For temporal modulation at 16kHz sample rate with 100 frames
/// let freqs_hz = fftfreq::<f64>(100, 1.0 / 16000.0);
/// // Returns frequencies in Hz
/// ```
#[inline]
#[must_use]
pub fn fftfreq<T: Sample>(n: usize, d: f64) -> Vec<T> {
    let mut freqs = Vec::with_capacity(n);
    let n_f64 = n as f64;
    let n_half = n.div_ceil(2);

    // Positive frequencies: 0, 1, 2, ..., (n-1)/2
    for i in 0..n_half {
        freqs.push(T::from_f64(i as f64 / (n_f64 * d)));
    }

    // Negative frequencies: -n/2, ..., -2, -1
    for i in n_half..n {
        freqs.push(T::from_f64((i as f64 - n_f64) / (n_f64 * d)));
    }

    freqs
}

/// Compute FFT sample frequencies for real FFT (like `numpy.fft.rfftfreq`).
///
/// Returns only the positive frequencies for a real-to-complex FFT.
/// For an FFT of length `n` with sample spacing `d`, returns `n/2 + 1` frequencies.
///
/// # Arguments
///
/// * `n` - Window length (number of samples in original real signal)
/// * `d` - Sample spacing (inverse of sampling rate). Default is 1.0.
///
/// # Returns
///
/// Vector of length `n/2 + 1` containing the positive frequency bin centers.
///
/// # Examples
///
/// ```rust
/// use spectrograms::fft2d::rfftfreq;
///
/// // For 8 samples
/// let freqs = rfftfreq::<f64>(8, 1.0);
/// // Returns: [0.0, 0.125, 0.25, 0.375, 0.5]
/// ```
#[inline]
#[must_use]
pub fn rfftfreq<T: Sample>(n: usize, d: f64) -> Vec<T> {
    let n_out = n / 2 + 1;
    let mut freqs = Vec::with_capacity(n_out);
    let n_f64 = n as f64;

    for i in 0..n_out {
        freqs.push(T::from_f64(i as f64 / (n_f64 * d)));
    }

    freqs
}

/// 2D FFT planner for efficient batch processing.
///
/// Caches FFT plans internally, avoiding repeated setup overhead when
/// processing multiple arrays with the same dimensions.
///
/// # Examples
///
/// ```rust
/// use spectrograms::fft2d::Fft2dPlanner;
/// use ndarray::Array2;
///
/// let mut planner = Fft2dPlanner::new();
///
/// // Process multiple images with same size
/// for _ in 0..10 {
///     let img = Array2::<f64>::zeros((128, 128));
///     let spectrum = planner.fft2d(&img.view()).unwrap();
/// }
/// ```
pub struct Fft2dPlanner<T: Sample = f64> {
    cache_r2c: HashMap<(usize, usize), T::R2cPlan2d>,
    cache_c2r: HashMap<(usize, usize), T::C2rPlan2d>,
}

impl<T: Sample> Fft2dPlanner<T> {
    /// Create a new 2D FFT planner.
    ///
    /// # Returns
    ///
    /// A new `Fft2dPlanner` instance.
    #[inline]
    #[must_use]
    pub fn new() -> Self {
        Self {
            cache_r2c: HashMap::new(),
            cache_c2r: HashMap::new(),
        }
    }

    /// Compute 2D FFT using cached plans.
    ///
    /// Reuses plans for arrays with dimensions seen before, avoiding setup overhead.
    ///
    /// # Arguments
    ///
    /// * `data` - Input 2D array with shape `(nrows, ncols)`
    ///
    /// # Returns
    ///
    /// Complex 2D array with shape `(nrows, ncols/2 + 1)` containing the frequency domain representation.
    ///
    /// # Errors
    ///
    /// Returns `SpectrogramError` if the input array is invalid or processing fails.
    #[inline]
    pub fn fft2d(&mut self, data: &ArrayView2<T>) -> SpectrogramResult<Array2<Complex<T>>> {
        let (nrows, ncols) = (data.nrows(), data.ncols());

        if nrows == 0 || ncols == 0 {
            return Err(SpectrogramError::invalid_input(
                "array dimensions must be > 0",
            ));
        }

        if !data.is_standard_layout() {
            return Err(SpectrogramError::invalid_input(
                "array must be contiguous and row-major (standard layout)",
            ));
        }

        let out_shape = r2c_output_size_2d(nrows, ncols);
        let plan = match self.cache_r2c.entry((nrows, ncols)) {
            Entry::Occupied(e) => e.into_mut(),
            Entry::Vacant(e) => e.insert(T::plan_r2c_2d(nrows, ncols)?),
        };

        let input_slice = data
            .as_slice()
            .ok_or_else(|| SpectrogramError::invalid_input("array must be contiguous"))?;

        let mut output = vec![Complex::new(T::zero(), T::zero()); out_shape.0 * out_shape.1];
        plan.process(input_slice, &mut output)?;

        Array2::from_shape_vec(out_shape, output)
            .map_err(|e| SpectrogramError::invalid_input(format!("failed to reshape output: {e}")))
    }

    /// Compute inverse 2D FFT using cached plans.
    ///
    /// # Arguments
    ///
    /// * `spectrum` - Complex frequency array with shape `(nrows, ncols/2 + 1)`
    /// * `output_ncols` - Number of columns in the output (must match original)
    ///
    /// # Returns
    ///
    /// Real 2D array with shape `(nrows, output_ncols)`.
    ///
    /// # Errors
    ///
    /// Returns `SpectrogramError` if the input array is invalid or processing fails.    
    #[inline]
    pub fn ifft2d(
        &mut self,
        spectrum: &ArrayView2<Complex<T>>,
        output_ncols: usize,
    ) -> SpectrogramResult<Array2<T>> {
        let nrows = spectrum.nrows();

        if nrows == 0 || output_ncols == 0 {
            return Err(SpectrogramError::invalid_input("dimensions must be > 0"));
        }

        let expected_ncols = output_ncols / 2 + 1;
        if spectrum.ncols() != expected_ncols {
            return Err(SpectrogramError::dimension_mismatch(
                expected_ncols,
                spectrum.ncols(),
            ));
        }

        if !spectrum.is_standard_layout() {
            return Err(SpectrogramError::invalid_input(
                "array must be contiguous and row-major (standard layout)",
            ));
        }

        let plan = match self.cache_c2r.entry((nrows, output_ncols)) {
            Entry::Occupied(e) => e.into_mut(),
            Entry::Vacant(e) => e.insert(T::plan_c2r_2d(nrows, output_ncols)?),
        };

        let input_slice = spectrum
            .as_slice()
            .ok_or_else(|| SpectrogramError::invalid_input("array must be contiguous"))?;

        let mut output = vec![T::zero(); nrows * output_ncols];
        plan.process(input_slice, &mut output)?;

        Array2::from_shape_vec((nrows, output_ncols), output)
            .map_err(|e| SpectrogramError::invalid_input(format!("failed to reshape output: {e}")))
    }

    /// Compute 2D power spectrum using cached plans.
    ///
    /// # Arguments
    ///
    /// * `data` - Input 2D array with shape `(nrows, ncols)`
    ///
    /// # Returns
    ///
    /// Real 2D array with shape `(nrows, ncols/2 + 1)` containing the power spectrum.
    ///
    /// # Errors
    ///
    /// Returns `SpectrogramError` if the input array is invalid or processing fails.
    #[inline]
    pub fn power_spectrum_2d(&mut self, data: &ArrayView2<T>) -> SpectrogramResult<Array2<T>> {
        let spectrum = self.fft2d(data)?;
        let power = spectrum.mapv(|c| c.norm_sqr());
        Ok(power)
    }

    /// Compute 2D magnitude spectrum using cached plans.
    ///
    /// # Arguments
    ///
    /// * `data` - Input 2D array with shape `(nrows, ncols)`
    ///
    /// # Returns
    ///
    /// Real 2D array with shape `(nrows, ncols/2 + 1)` containing the magnitude spectrum.
    ///
    /// # Errors
    ///
    /// Returns `SpectrogramError` if the input array is invalid or processing fails.
    #[inline]
    pub fn magnitude_spectrum_2d(&mut self, data: &ArrayView2<T>) -> SpectrogramResult<Array2<T>> {
        let spectrum = self.fft2d(data)?;
        let magnitude = spectrum.mapv(num_complex::Complex::norm);
        Ok(magnitude)
    }
}

impl<T: Sample> Default for Fft2dPlanner<T> {
    #[inline]
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn test_fft2d_zeros() {
        let data = Array2::<f64>::zeros((32, 32));
        let result = fft2d(&data.view());
        assert!(result.is_ok());
        let spectrum = result.unwrap();
        assert_eq!(spectrum.shape(), &[32, 17]); // 32 x (32/2 + 1)

        // All values should be ~zero
        for val in spectrum.iter() {
            assert!(val.norm() < 1e-10);
        }
    }

    #[test]
    fn test_fft2d_ones() {
        let data = Array2::<f64>::ones((32, 32));
        let spectrum = fft2d(&data.view()).unwrap();

        // DC component should have all energy
        assert!(spectrum[[0, 0]].norm() > 1000.0);

        // Other components should be ~zero
        for i in 1..32 {
            assert!(spectrum[[i, 0]].norm() < 1e-10);
        }
    }

    #[test]
    fn test_roundtrip() {
        let original = Array2::<f64>::from_shape_fn((64, 64), |(i, j)| {
            (i as f64 * 0.1).sin() + (j as f64 * 0.2).cos()
        });

        let spectrum = fft2d(&original.view()).unwrap();
        let reconstructed = ifft2d(&spectrum, 64).unwrap();

        for ((i, j), &val) in original.indexed_iter() {
            assert!((reconstructed[[i, j]] - val).abs() < 1e-10);
        }
    }

    #[test]
    fn test_planner_reuse() {
        let mut planner = Fft2dPlanner::new();

        for _ in 0..5 {
            let data = Array2::<f64>::zeros((32, 32));
            let result = planner.fft2d(&data.view());
            assert!(result.is_ok());
        }
    }
}