veils 0.1.4

// Portions of this implementation are adapted from SciPy's
// scipy/signal/_short_time_fft.py. See LICENSE-SCIPY.
//! # Veils: High-Performance STFT Library
//!
//! Veils provides a high-performance Short-Time Fourier Transform (STFT) implementation
//! compatible with the one-dimensional, real-valued subset of
//! `scipy.signal.ShortTimeFFT` documented in the project README.
//!
//! ## Features
//!
//! - **SciPy compatibility**: Verified against SciPy 1.18.1
//! - **High performance**: Optimized Rust implementation using rustfft
//! - **Multiple FFT modes**: TwoSided, Centered, OneSided, OneSided2X
//! - **Invertible STFT**: Perfect reconstruction with canonical dual window
//! - **Flexible windowing**: Support for custom windows and dual windows
//!
//! ## Quick Start
//!
//! ```rust
//! use veils::StandaloneSTFT;
//!
//! // Create a Hann window
//! let window: Vec<f64> = (0..16)
//!     .map(|i| 0.5 * (1.0 - (2.0 * std::f64::consts::PI * i as f64 / 15.0).cos()))
//!     .collect();
//!
//! // Create STFT instance
//! let stft = StandaloneSTFT::new(
//!     window,
//!     4,      // hop length
//!     1000.0, // sampling rate
//!     Some("onesided"),
//!     None,   // mfft (defaults to window length)
//!     None,   // dual_win (computed automatically)
//!     None,   // phase_shift
//! ).unwrap();
//!
//! // Generate test signal
//! let signal: Vec<f64> = (0..64)
//!     .map(|i| (2.0 * std::f64::consts::PI * 5.0 * i as f64 / 1000.0).sin())
//!     .collect();
//!
//! // Compute STFT
//! let stft_result = stft.stft(&signal, None, None, None).unwrap();
//! println!("STFT shape: {} x {}", stft_result.len(), stft_result[0].len());
//! ```

use num_complex::Complex;
use rustfft::{Fft, FftPlanner};
use std::str::FromStr;
use std::sync::Arc;

/// FFT mode for STFT computation
///
/// Determines how the FFT is computed and what frequency range is returned.
#[derive(Debug, Clone)]
pub enum FftMode {
    /// Two-sided FFT: returns full frequency spectrum
    TwoSided,
    /// Centered FFT: two-sided with fftshift applied
    Centered,
    /// One-sided FFT: returns only positive frequencies (for real signals)
    OneSided,
    /// One-sided FFT with 2x scaling: like OneSided but with amplitude scaling
    OneSided2X,
}

impl FromStr for FftMode {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "twosided" => Ok(FftMode::TwoSided),
            "centered" => Ok(FftMode::Centered),
            "onesided" => Ok(FftMode::OneSided),
            "onesided2X" => Ok(FftMode::OneSided2X),
            _ => Err(format!("Unknown fft_mode: {}", s)),
        }
    }
}

impl FftMode {
    pub fn is_onesided(&self) -> bool {
        matches!(self, FftMode::OneSided | FftMode::OneSided2X)
    }
}

/// Short-Time Fourier Transform implementation
///
/// This struct provides a real-valued STFT implementation compatible with the
/// supported subset of SciPy's `ShortTimeFFT`.
/// It supports forward STFT, inverse STFT (ISTFT), and various FFT modes.
///
/// # Examples
///
/// ```rust,no_run
/// use veils::StandaloneSTFT;
///
/// // Create a simple Hann window
/// let window: Vec<f64> = (0..16)
///     .map(|i| 0.5 * (1.0 - (2.0 * std::f64::consts::PI * i as f64 / 15.0).cos()))
///     .collect();
///
/// let mut stft = StandaloneSTFT::new(
///     window, 4, 1000.0, Some("onesided"), None, None, None
/// ).unwrap();
///
/// // Test signal: sine wave
/// let signal: Vec<f64> = (0..64)
///     .map(|i| (2.0 * std::f64::consts::PI * 5.0 * i as f64 / 1000.0).sin())
///     .collect();
///
/// // Forward STFT
/// let stft_result = stft.stft(&signal, None, None, None).unwrap();
///
/// // Inverse STFT (perfect reconstruction)
/// let reconstructed = stft.istft(&stft_result, None, None).unwrap();
/// ```
pub struct StandaloneSTFT {
    win: Vec<Complex<f64>>,
    hop: usize,
    fs: f64,
    fft_mode: FftMode,
    mfft: usize,
    phase_shift: i32,
    dual_win: Option<Vec<Complex<f64>>>,
    // FFT planners
    forward_fft: Option<Arc<dyn Fft<f64>>>,
    inverse_fft: Option<Arc<dyn Fft<f64>>>,
}

impl StandaloneSTFT {
    /// Create a new STFT instance
    ///
    /// # Arguments
    ///
    /// * `win` - Window function as a vector of f64 values
    /// * `hop` - Hop length (number of samples between adjacent STFT columns)
    /// * `fs` - Sampling frequency in Hz
    /// * `fft_mode` - FFT mode: "twosided", "centered", "onesided", or "onesided2X"
    /// * `mfft` - FFT length (defaults to window length if None)
    /// * `dual_win` - Dual window for ISTFT (computed automatically if None)
    /// * `phase_shift` - Phase shift for FFT (default: 0)
    ///
    /// # Returns
    ///
    /// Returns `Ok(StandaloneSTFT)` on success, or `Err(String)` with error message.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use veils::StandaloneSTFT;
    ///
    /// // Hann window of length 16
    /// let window: Vec<f64> = (0..16)
    ///     .map(|i| 0.5 * (1.0 - (2.0 * std::f64::consts::PI * i as f64 / 15.0).cos()))
    ///     .collect();
    ///
    /// let stft = StandaloneSTFT::new(
    ///     window,
    ///     4,      // hop length
    ///     1000.0, // sampling rate
    ///     Some("onesided"),
    ///     None,   // use window length for FFT
    ///     None,   // compute dual window automatically
    ///     None,   // no phase shift
    /// ).unwrap();
    /// ```
    pub fn new(
        win: Vec<f64>,
        hop: usize,
        fs: f64,
        fft_mode: Option<&str>,
        mfft: Option<usize>,
        dual_win: Option<Vec<f64>>,
        phase_shift: Option<i32>,
    ) -> Result<Self, String> {
        // Validate inputs
        if win.is_empty() {
            return Err("Parameter win must be non-empty!".to_string());
        }

        if !win.iter().all(|&x| x.is_finite()) {
            return Err("Parameter win must have finite entries!".to_string());
        }

        if hop < 1 {
            return Err(format!("Parameter hop={} is not >= 1!", hop));
        }

        if fs <= 0.0 || fs.is_nan() {
            return Err(format!("Sampling frequency fs={} must be positive!", fs));
        }

        // Convert window to complex
        let win_complex: Vec<Complex<f64>> =
            win.into_iter().map(|x| Complex::new(x, 0.0)).collect();

        let fft_mode = fft_mode.unwrap_or("onesided").parse()?;
        let mfft = mfft.unwrap_or(win_complex.len());
        let phase_shift = phase_shift.unwrap_or(0);

        if mfft < win_complex.len() {
            return Err(format!(
                "mfft={} must be at least the window length {}!",
                mfft,
                win_complex.len()
            ));
        }
        if !(-(mfft as i64) < phase_shift as i64 && (phase_shift as i64) < mfft as i64) {
            return Err(format!(
                "-mfft < phase_shift < mfft does not hold for mfft={}, phase_shift={}!",
                mfft, phase_shift
            ));
        }

        // Convert dual window if provided
        let dual_win_complex = if let Some(dw) = dual_win {
            if dw.len() != win_complex.len() {
                return Err("dual_win.len() must equal win.len()!".to_string());
            }
            if !dw.iter().all(|&x| x.is_finite()) {
                return Err("Parameter dual_win must have finite entries!".to_string());
            }
            Some(dw.into_iter().map(|x| Complex::new(x, 0.0)).collect())
        } else {
            None
        };

        let mut planner = FftPlanner::new();
        let forward_fft = Some(planner.plan_fft_forward(mfft));
        let inverse_fft = Some(planner.plan_fft_inverse(mfft));

        Ok(StandaloneSTFT {
            win: win_complex,
            hop,
            fs,
            fft_mode,
            mfft,
            phase_shift,
            dual_win: dual_win_complex,
            forward_fft,
            inverse_fft,
        })
    }

    // Properties
    pub fn win(&self) -> &[Complex<f64>] {
        &self.win
    }

    pub fn hop(&self) -> usize {
        self.hop
    }

    pub fn fs(&self) -> f64 {
        self.fs
    }

    pub fn fft_mode(&self) -> &FftMode {
        &self.fft_mode
    }

    pub fn mfft(&self) -> usize {
        self.mfft
    }

    pub fn phase_shift(&self) -> i32 {
        self.phase_shift
    }

    pub fn sampling_period(&self) -> f64 {
        1.0 / self.fs
    }

    pub fn delta_t(&self) -> f64 {
        self.hop as f64 * self.sampling_period()
    }

    pub fn delta_f(&self) -> f64 {
        1.0 / (self.mfft as f64 * self.sampling_period())
    }

    pub fn m_num(&self) -> usize {
        self.win.len()
    }

    pub fn m_num_mid(&self) -> usize {
        self.m_num() / 2
    }

    pub fn f_pts(&self) -> usize {
        if self.fft_mode.is_onesided() {
            self.mfft / 2 + 1
        } else {
            self.mfft
        }
    }

    pub fn onesided_fft(&self) -> bool {
        self.fft_mode.is_onesided()
    }

    /// Calculate canonical dual window for 1d window and a time step of hop samples
    fn calc_dual_canonical_window(&self) -> Result<Vec<Complex<f64>>, String> {
        let win = &self.win;
        let hop = self.hop;

        if hop > win.len() {
            return Err(format!(
                "hop={} is larger than window length of {} => STFT not invertible!",
                hop,
                win.len()
            ));
        }

        // w2 = win.real**2 + win.imag**2
        let w2: Vec<f64> = win.iter().map(|w| w.norm_sqr()).collect();
        let mut dd = w2.clone();

        // Overlap-add for dual window calculation
        for k in (hop..win.len()).step_by(hop) {
            // DD[k_:] += w2[:-k_]
            for i in k..win.len() {
                dd[i] += w2[i - k];
            }
            // DD[:-k_] += w2[k_:]
            for i in 0..(win.len() - k) {
                dd[i] += w2[i + k];
            }
        }

        // Check DD > 0 (using relative resolution)
        let max_dd = dd.iter().fold(0.0f64, |a, &b| a.max(b));
        let relative_resolution = f64::EPSILON * max_dd;

        if !dd.iter().all(|&x| x >= relative_resolution) {
            return Err("Short-time Fourier Transform not invertible!".to_string());
        }

        // Return win / DD
        Ok(win.iter().zip(dd.iter()).map(|(w, &d)| w / d).collect())
    }

    pub fn dual_win(&mut self) -> Result<&[Complex<f64>], String> {
        if self.dual_win.is_none() {
            self.dual_win = Some(self.calc_dual_canonical_window()?);
        }
        Ok(self.dual_win.as_ref().unwrap())
    }

    pub fn invertible(&mut self) -> bool {
        self.dual_win().is_ok()
    }

    /// Apply FFT based on fft_mode
    fn fft_func(&self, mut x: Vec<Complex<f64>>) -> Vec<Complex<f64>> {
        // Handle phase shift first (like scipy does) - always apply if phase_shift is set
        // Note: phase_shift is always Some(0) by default, so we always apply this
        if x.len() < self.mfft {
            // Zero pad if needed
            x.resize(self.mfft, Complex::new(0.0, 0.0));
        }
        let p_s =
            (self.phase_shift + self.m_num_mid() as i32).rem_euclid(self.m_num() as i32) as usize;
        // Equivalent to np.roll(x, -p_s, axis=-1)
        x.rotate_left(p_s);

        // Ensure we have the right size for FFT
        if x.len() != self.mfft {
            x.resize(self.mfft, Complex::new(0.0, 0.0));
        }

        // Compute FFT based on mode
        match &self.fft_mode {
            FftMode::TwoSided => {
                if let Some(ref fft) = self.forward_fft {
                    fft.process(&mut x);
                }
                x
            }
            FftMode::Centered => {
                if let Some(ref fft) = self.forward_fft {
                    fft.process(&mut x);
                }
                // Apply fftshift
                self.fftshift(x)
            }
            FftMode::OneSided => {
                // For real FFT, we should only use the real part of input
                // and compute a full FFT, then take the first half + 1
                // But since rustfft doesn't have a dedicated real FFT, we simulate it
                if let Some(ref fft) = self.forward_fft {
                    fft.process(&mut x);
                }
                let n_out = self.mfft / 2 + 1;
                x.truncate(n_out);

                // For real input signals, the result should have conjugate symmetry
                // The DC and Nyquist components should be real
                if n_out > 0 {
                    x[0] = Complex::new(x[0].re, 0.0); // DC component
                }
                if self.mfft.is_multiple_of(2) && n_out > 1 {
                    x[n_out - 1] = Complex::new(x[n_out - 1].re, 0.0); // Nyquist component
                }
                x
            }
            FftMode::OneSided2X => {
                // Similar to OneSided but with scaling
                if let Some(ref fft) = self.forward_fft {
                    fft.process(&mut x);
                }
                let n_out = self.mfft / 2 + 1;
                x.truncate(n_out);

                // Apply scaling (factor of 2 for unpaired frequencies)
                let fac = 2.0; // Assuming no PSD scaling for now
                if self.mfft.is_multiple_of(2) {
                    // For even input length, the last entry is unpaired
                    let len = x.len();
                    for x_i in x.iter_mut().skip(1).take(len - 2) {
                        *x_i *= fac;
                    }
                } else {
                    for x_i in x.iter_mut().skip(1) {
                        *x_i *= fac;
                    }
                }
                x
            }
        }
    }

    /// Apply IFFT based on fft_mode
    fn ifft_func(&self, mut x: Vec<Complex<f64>>) -> Vec<Complex<f64>> {
        match &self.fft_mode {
            FftMode::TwoSided => {
                if let Some(ref ifft) = self.inverse_fft {
                    ifft.process(&mut x);
                }
            }
            FftMode::Centered => {
                // Apply ifftshift first
                x = self.ifftshift(x);
                if let Some(ref ifft) = self.inverse_fft {
                    ifft.process(&mut x);
                }
            }
            FftMode::OneSided => {
                // Reconstruct full spectrum from one-sided
                let mut full_x = vec![Complex::new(0.0, 0.0); self.mfft];
                full_x[..x.len()].copy_from_slice(&x);

                // Mirror the spectrum (conjugate symmetry)
                // Skip DC (index 0) and Nyquist frequency (last element for even mfft)
                let mirror_end = if self.mfft.is_multiple_of(2) {
                    // For even mfft, don't mirror the Nyquist frequency
                    (self.mfft / 2).min(x.len() - 1)
                } else {
                    // For odd mfft, mirror all except DC
                    (self.mfft / 2).min(x.len() - 1)
                };

                for i in 1..=mirror_end {
                    if i < x.len() && (self.mfft - i) < full_x.len() {
                        full_x[self.mfft - i] = x[i].conj();
                    }
                }

                if let Some(ref ifft) = self.inverse_fft {
                    ifft.process(&mut full_x);
                }
                x = full_x;
            }
            FftMode::OneSided2X => {
                // Undo scaling first
                let mut xc = x.clone();
                let fac = 2.0;

                if self.mfft.is_multiple_of(2) {
                    let len = xc.len();
                    for xc_i in xc.iter_mut().skip(1).take(len - 2) {
                        *xc_i /= fac;
                    }
                } else {
                    for xc_i in xc.iter_mut().skip(1) {
                        *xc_i /= fac;
                    }
                }

                // Reconstruct full spectrum
                let mut full_x = vec![Complex::new(0.0, 0.0); self.mfft];
                full_x[..xc.len()].copy_from_slice(&xc);

                // Mirror the spectrum (conjugate symmetry) - same fix as OneSided
                let mirror_end = (self.mfft / 2).min(xc.len() - 1);

                for i in 1..=mirror_end {
                    if i < xc.len() && (self.mfft - i) < full_x.len() {
                        full_x[self.mfft - i] = xc[i].conj();
                    }
                }

                if let Some(ref ifft) = self.inverse_fft {
                    ifft.process(&mut full_x);
                }
                x = full_x;
            }
        }

        // CRITICAL FIX: Apply correct normalization to match scipy behavior
        // RustFFT produces unnormalized results, scipy.fft.irfft produces normalized results
        // For onesided IFFT, we need to normalize by the correct factor to match scipy
        let normalization_factor = 1.0 / (self.mfft as f64);
        for val in &mut x {
            *val *= normalization_factor;
        }

        // Handle phase shift - always apply like scipy does
        let p_s =
            (self.phase_shift + self.m_num_mid() as i32).rem_euclid(self.m_num() as i32) as usize;
        x.rotate_right(p_s);

        // Return only the window length
        x.truncate(self.m_num());
        x
    }

    /// FFT shift implementation
    fn fftshift(&self, mut x: Vec<Complex<f64>>) -> Vec<Complex<f64>> {
        let n = x.len();
        let mid = n.div_ceil(2);
        x.rotate_left(mid);
        x
    }

    /// Inverse FFT shift implementation
    fn ifftshift(&self, mut x: Vec<Complex<f64>>) -> Vec<Complex<f64>> {
        let n = x.len();
        let mid = n / 2;
        x.rotate_left(mid);
        x
    }

    /// Calculate padding parameters
    fn pre_padding(&self) -> (i32, i32) {
        let w2: Vec<f64> = self.win.iter().map(|w| w.norm_sqr()).collect();
        let n0 = -(self.m_num_mid() as i32);

        // Python: for p_, n_ in enumerate(range(n0, n0-self.m_num-1, -self.hop)):
        let mut p = 0;
        let mut n = n0;
        loop {
            let n_next = n - self.hop as i32;
            if n_next + self.m_num() as i32 <= 0 {
                return (n, -p);
            }
            // Check if n_next is valid and all w2[n_next:] == 0
            let suffix_start = if n_next < 0 {
                self.m_num() as i32 + n_next
            } else {
                n_next
            };
            if suffix_start >= 0
                && (suffix_start as usize) < w2.len()
                && w2[(suffix_start as usize)..].iter().all(|&x| x == 0.0)
            {
                return (n, -p);
            }

            n = n_next;
            p += 1;

            // Break condition to match Python range
            if n < n0 - self.m_num() as i32 {
                break;
            }
        }

        // Fallback
        (
            n0 - self.m_num() as i32,
            -((self.m_num() / self.hop + 1) as i32),
        )
    }

    fn post_padding(&self, n: usize) -> (i32, i32) {
        let m2p = self.m_num() - self.m_num_mid();
        if n < m2p {
            panic!("Parameter n must be >= ceil(m_num/2) = {}!", m2p);
        }

        let w2: Vec<f64> = self.win.iter().map(|w| w.norm_sqr()).collect();
        let q1 = n / self.hop;
        let k1 = q1 as i32 * self.hop as i32 - self.m_num_mid() as i32;

        for (q_offset, k) in (k1..(n as i32 + self.m_num() as i32))
            .step_by(self.hop)
            .enumerate()
        {
            let q = q1 + q_offset;
            let n_next = k + self.hop as i32;
            if n_next >= n as i32 || w2[..(n as i32 - n_next) as usize].iter().all(|&x| x == 0.0) {
                return (k + self.m_num() as i32, q as i32 + 1);
            }
        }
        // Fallback
        (
            n as i32 + self.m_num() as i32,
            ((n + self.m_num_mid()) / self.hop + 1) as i32,
        )
    }

    pub fn p_min(&self) -> i32 {
        self.pre_padding().1
    }

    pub fn p_max(&self, n: usize) -> i32 {
        self.post_padding(n).1
    }

    pub fn p_range(
        &self,
        n: usize,
        p0: Option<i32>,
        p1: Option<i32>,
    ) -> Result<(i32, i32), String> {
        let p0 = p0.unwrap_or(self.p_min());
        let p1 = p1.unwrap_or(self.p_max(n));

        if !(self.p_min() <= p0 && p0 < p1 && p1 <= self.p_max(n)) {
            return Err(format!("Invalid slice range: p0={}, p1={}", p0, p1));
        }

        Ok((p0, p1))
    }

    /// Generate windowed slices of input signal
    fn x_slices(&self, x: &[f64], k_offset: i32, p0: i32, p1: i32) -> Vec<Vec<Complex<f64>>> {
        let n = x.len();
        let mut slices = Vec::new();

        for p in p0..p1 {
            let k_center = p * self.hop as i32 + k_offset;
            let k_start = k_center - self.m_num_mid() as i32;
            let k_end = k_start + self.m_num() as i32;

            let mut slice_data = vec![Complex::new(0.0, 0.0); self.m_num()];

            // Extract slice with padding if necessary
            if k_start >= 0 && k_end <= n as i32 {
                // No padding needed
                let start_idx = k_start as usize;
                let end_idx = k_end as usize;
                for (i, &val) in x[start_idx..end_idx].iter().enumerate() {
                    slice_data[i] = Complex::new(val, 0.0);
                }
            } else {
                // Padding needed - using zeros padding for now
                let valid_start = 0.max(k_start);
                let valid_end = (n as i32).min(k_end);

                if valid_start < valid_end {
                    let slice_start = valid_start - k_start;
                    let valid_start_idx = valid_start as usize;
                    let valid_end_idx = valid_end as usize;
                    let slice_start_idx = slice_start as usize;

                    for (i, &val) in x[valid_start_idx..valid_end_idx].iter().enumerate() {
                        slice_data[slice_start_idx + i] = Complex::new(val, 0.0);
                    }
                }
            }

            slices.push(slice_data);
        }

        slices
    }

    /// Perform the short-time Fourier transform
    ///
    /// Computes the STFT of the input signal using the configured window and parameters.
    ///
    /// # Arguments
    ///
    /// * `x` - Input signal as a slice of f64 values
    /// * `p0` - Start index for STFT computation (None for automatic)
    /// * `p1` - End index for STFT computation (None for automatic)
    /// * `k_offset` - Sample offset for time axis (default: 0)
    ///
    /// # Returns
    ///
    /// Returns `Ok(Vec<Vec<Complex<f64>>>)` where the outer vector represents time slices
    /// and the inner vector represents frequency bins, or `Err(String)` on error.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use veils::StandaloneSTFT;
    ///
    /// let window: Vec<f64> = (0..16)
    ///     .map(|i| 0.5 * (1.0 - (2.0 * std::f64::consts::PI * i as f64 / 15.0).cos()))
    ///     .collect();
    ///
    /// let stft = StandaloneSTFT::new(window, 4, 1000.0, Some("onesided"), None, None, None).unwrap();
    ///
    /// let signal: Vec<f64> = (0..64)
    ///     .map(|i| (2.0 * std::f64::consts::PI * 5.0 * i as f64 / 1000.0).sin())
    ///     .collect();
    ///
    /// let stft_result = stft.stft(&signal, None, None, None).unwrap();
    /// println!("STFT computed: {} time slices, {} frequency bins",
    ///          stft_result.len(), stft_result[0].len());
    /// ```
    pub fn stft(
        &self,
        x: &[f64],
        p0: Option<i32>,
        p1: Option<i32>,
        k_offset: Option<i32>,
    ) -> Result<Vec<Vec<Complex<f64>>>, String> {
        if self.onesided_fft() && x.iter().any(|&val| val.is_nan()) {
            return Err("Complex-valued input not allowed for one-sided FFT modes".to_string());
        }

        let n = x.len();
        let m2p = self.m_num() - self.m_num_mid();
        if n < m2p {
            return Err(format!(
                "Input length {} must be >= ceil(m_num/2) = {}",
                n, m2p
            ));
        }

        let (p0, p1) = self.p_range(n, p0, p1)?;
        let k_offset = k_offset.unwrap_or(0);

        let slices = self.x_slices(x, k_offset, p0, p1);
        let mut time_slice_results = Vec::new();

        for slice in slices.iter() {
            // Apply window and compute FFT
            let windowed: Vec<Complex<f64>> = slice
                .iter()
                .zip(self.win.iter())
                .map(|(x, w)| {
                    let result = x * w.conj();
                    // For one-sided FFT, ensure input is real
                    if self.onesided_fft() {
                        Complex::new(result.re, 0.0)
                    } else {
                        result
                    }
                })
                .collect();

            let fft_result = self.fft_func(windowed);
            time_slice_results.push(fft_result);
        }

        // Transpose to match scipy format: [frequency_bins][time_slices]
        let num_freq_bins = self.f_pts();
        let num_time_slices = time_slice_results.len();
        let mut stft_result = vec![vec![Complex::new(0.0, 0.0); num_time_slices]; num_freq_bins];

        for (time_idx, time_slice) in time_slice_results.iter().enumerate() {
            for (freq_idx, &freq_val) in time_slice.iter().enumerate() {
                if freq_idx < num_freq_bins {
                    stft_result[freq_idx][time_idx] = freq_val;
                }
            }
        }

        Ok(stft_result)
    }

    /// Inverse short-time Fourier transform (ISTFT)
    ///
    /// Reconstructs a time-domain signal from its STFT representation with perfect reconstruction.
    ///
    /// # Arguments
    ///
    /// * `stft_data` - STFT data in format \[frequency_bins\]\[time_slices\]
    /// * `k0` - Start sample for reconstruction (None for automatic)
    /// * `k1` - End sample for reconstruction (None for automatic)
    ///
    /// # Returns
    ///
    /// Returns `Ok(Vec<f64>)` with the reconstructed signal, or `Err(String)` on error.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use veils::StandaloneSTFT;
    ///
    /// let window: Vec<f64> = (0..16)
    ///     .map(|i| 0.5 * (1.0 - (2.0 * std::f64::consts::PI * i as f64 / 15.0).cos()))
    ///     .collect();
    ///
    /// let mut stft = StandaloneSTFT::new(window, 4, 1000.0, Some("onesided"), None, None, None).unwrap();
    ///
    /// let signal: Vec<f64> = (0..64)
    ///     .map(|i| (2.0 * std::f64::consts::PI * 5.0 * i as f64 / 1000.0).sin())
    ///     .collect();
    ///
    /// // Forward STFT
    /// let stft_result = stft.stft(&signal, None, None, None).unwrap();
    ///
    /// // Inverse STFT - perfect reconstruction
    /// let reconstructed = stft.istft(&stft_result, None, None).unwrap();
    ///
    /// // Verify reconstruction accuracy
    /// let error: f64 = signal.iter().zip(reconstructed.iter())
    ///     .map(|(a, b)| (a - b).abs())
    ///     .fold(0.0, f64::max);
    /// assert!(error < 1e-10);
    /// ```
    pub fn istft(
        &mut self,
        stft_data: &[Vec<Complex<f64>>],
        k0: Option<i32>,
        k1: Option<i32>,
    ) -> Result<Vec<f64>, String> {
        if stft_data.is_empty() {
            return Err("STFT data cannot be empty".to_string());
        }

        // CRITICAL: stft_data is in Python format [freq][time], we need to check dimensions correctly
        let f_pts_expected = self.f_pts();
        let time_slices = stft_data[0].len();

        if stft_data.len() != f_pts_expected {
            return Err(format!(
                "STFT frequency dimension {} must equal {}",
                stft_data.len(),
                f_pts_expected
            ));
        }
        if stft_data.iter().any(|row| row.len() != time_slices) {
            return Err("STFT data must have equal-length frequency rows".to_string());
        }

        let n_min = self.m_num() - self.m_num_mid();
        let q_num = self.p_max(n_min) - self.p_min();
        if (time_slices as i32) < q_num {
            return Err(format!(
                "STFT time dimension {} needs to have at least {} slices",
                time_slices, q_num
            ));
        }

        let q_max = time_slices as i32 + self.p_min();
        let k_max = (q_max - 1) * self.hop as i32 + self.m_num() as i32 - self.m_num_mid() as i32;

        let k0 = k0.unwrap_or(0);
        let k1 = k1.unwrap_or(k_max);

        let k_min = self.pre_padding().0;
        if !(k_min <= k0 && k0 < k1 && k1 <= k_max) {
            return Err(format!(
                "k_min={} <= k0={} < k1={} <= k_max={} does not hold",
                k_min, k0, k1, k_max
            ));
        }

        let num_pts = (k1 - k0) as usize;
        if num_pts < n_min {
            return Err(format!(
                "Output length {} has to be at least {}",
                num_pts, n_min
            ));
        }

        // Match Python's q0 calculation exactly
        let q0 = if k0 >= 0 {
            k0.div_euclid(self.hop as i32) + self.p_min()
        } else {
            k0.div_euclid(self.hop as i32)
        };
        let q1 = (self.p_max(k1 as usize)).min(q_max);

        let dual_win = self.dual_win()?.to_vec();
        let mut x = vec![0.0f64; num_pts];

        // CRITICAL FIX: Match Python's ISTFT logic exactly
        for q in q0..q1 {
            let stft_idx = q - self.p_min();
            if stft_idx < 0 || stft_idx >= time_slices as i32 {
                continue; // No more STFT data available
            }

            // CRITICAL: Extract frequency slice like Python: S[:, q_ - self.p_min]
            let mut stft_slice = vec![Complex::new(0.0, 0.0); f_pts_expected];
            for f in 0..f_pts_expected {
                stft_slice[f] = stft_data[f][stft_idx as usize];
            }

            // Apply IFFT to get time domain signal
            let xs_raw = self.ifft_func(stft_slice);

            // Apply dual window (Python: xs = self._ifft_func(S[:, q_ - self.p_min]) * self.dual_win)
            let xs: Vec<Complex<f64>> = xs_raw
                .iter()
                .zip(dual_win.iter())
                .map(|(x, w)| x * w)
                .collect();

            // Calculate indices exactly like Python
            let i0 = q * self.hop as i32 - self.m_num_mid() as i32;
            let i1 = (i0 + self.m_num() as i32).min(num_pts as i32 + k0);
            let mut j0 = 0i32;
            let mut j1 = i1 - i0;

            let mut actual_i0 = i0;
            if i0 < k0 {
                // xs sticks out to the left on x
                j0 += k0 - i0;
                actual_i0 = k0;
            }

            if i1 > k0 + num_pts as i32 {
                // xs sticks out to the right
                j1 -= i1 - k0 - num_pts as i32;
            }

            // Apply overlap-add exactly like Python
            if actual_i0 < i1 && j0 < j1 {
                let target_start = (actual_i0 - k0) as usize;
                let _target_end = target_start + (j1 - j0) as usize;
                let source_start = j0 as usize;
                let source_end = j1 as usize;

                for (i, &val) in xs[source_start..source_end].iter().enumerate() {
                    if target_start + i < x.len() {
                        // Python: x[i0-k0:i1-k0] += xs[j0:j1].real if self.onesided_fft else xs[j0:j1]
                        if self.onesided_fft() {
                            x[target_start + i] += val.re;
                        } else {
                            x[target_start + i] += val.re; // For complex case, should be val itself, but we're returning f64
                        }
                    }
                }
            }
        }

        Ok(x)
    }

    pub fn t(
        &self,
        n: usize,
        p0: Option<i32>,
        p1: Option<i32>,
        k_offset: Option<i32>,
    ) -> Result<Vec<f64>, String> {
        let (p0, p1) = self.p_range(n, p0, p1)?;
        let k_offset = k_offset.unwrap_or(0);
        Ok((p0..p1)
            .map(|p| (p * self.hop as i32 + k_offset) as f64 * self.sampling_period())
            .collect())
    }

    pub fn f(&self) -> Vec<f64> {
        let mut freqs: Vec<f64> = if self.fft_mode.is_onesided() {
            (0..self.f_pts())
                .map(|i| i as f64 / (self.mfft as f64 * self.sampling_period()))
                .collect()
        } else {
            (0..self.mfft)
                .map(|i| {
                    let freq = i as f64 / (self.mfft as f64 * self.sampling_period());
                    if i >= self.mfft.div_ceil(2) {
                        freq - 1.0 / self.sampling_period()
                    } else {
                        freq
                    }
                })
                .collect()
        };

        if matches!(self.fft_mode, FftMode::Centered) {
            let shift = self.mfft.div_ceil(2);
            freqs.rotate_left(shift);
        }
        freqs
    }

    /// Debug method to access IFFT function
    pub fn debug_ifft_func(&self, x: Vec<Complex<f64>>) -> Vec<Complex<f64>> {
        self.ifft_func(x)
    }
}