mel_spec 0.5.0

//! Kaldi-style filterbank feature extraction.
//!
//! This module provides filterbank features matching kaldi_native_fbank output.
//!
//! # Parameters (matching kaldi defaults)
//! - `sample_rate`: 16000 Hz
//! - `num_mel_bins`: 80
//! - `frame_length_ms`: 25.0 ms (400 samples at 16kHz)
//! - `frame_shift_ms`: 10.0 ms (160 samples at 16kHz)
//! - `window_type`: povey (like hamming but goes to zero at edges)
//! - `dither`: 0.0 (disabled for inference)
//! - `preemphasis`: 0.97
//! - `remove_dc_offset`: true

use ndarray::{Array2, ArrayView1};
use realfft::{RealFftPlanner, RealToComplex};
use std::f64::consts::PI;
use std::sync::Arc;

use crate::mel::{SparseMelFilterbank, FRAME_LANES};

/// Configuration for Kaldi-compatible filterbank extraction.
#[derive(Clone, Debug)]
pub struct FbankConfig {
    pub sample_rate: f64,
    pub num_mel_bins: usize,
    pub frame_length_ms: f64,
    pub frame_shift_ms: f64,
    pub dither: f64,
    /// Energy floor for log computation (kaldi default: 0.0, uses FLT_EPSILON internally)
    pub energy_floor: f64,
    pub use_energy: bool,
    pub use_log_fbank: bool,
    pub use_power: bool,
    /// Preemphasis coefficient (kaldi default: 0.97)
    pub preemphasis: f64,
    /// If true, apply CMN (subtract mean across time for each frequency bin).
    pub apply_cmn: bool,
    /// Low frequency cutoff for mel filterbank (kaldi default: 20.0)
    pub low_freq: f64,
    /// High frequency cutoff (0.0 means Nyquist)
    pub high_freq: f64,
}

impl Default for FbankConfig {
    fn default() -> Self {
        Self {
            sample_rate: 16000.0,
            num_mel_bins: 80,
            frame_length_ms: 25.0,
            frame_shift_ms: 10.0,
            dither: 0.0,
            energy_floor: 0.0, // kaldi default
            use_energy: false,
            use_log_fbank: true,
            use_power: true,
            preemphasis: 0.97,
            apply_cmn: true,
            low_freq: 20.0,
            high_freq: 0.0, // 0 means Nyquist
        }
    }
}

impl FbankConfig {
    /// Frame length in samples.
    pub fn frame_length_samples(&self) -> usize {
        ((self.frame_length_ms / 1000.0) * self.sample_rate).round() as usize
    }

    /// Frame shift in samples.
    pub fn frame_shift_samples(&self) -> usize {
        ((self.frame_shift_ms / 1000.0) * self.sample_rate).round() as usize
    }

    /// Padded FFT size (next power of 2).
    pub fn fft_size(&self) -> usize {
        let frame_len = self.frame_length_samples();
        frame_len.next_power_of_two()
    }
}

/// Kaldi-compatible filterbank feature extractor.
pub struct Fbank {
    config: FbankConfig,
    mel_filters: Array2<f64>,
    sparse_mel_filters: SparseMelFilterbank,
    fft: Arc<dyn RealToComplex<f64>>,
    window: Vec<f64>,
}

impl Fbank {
    pub fn new(config: FbankConfig) -> Self {
        let fft_size = config.fft_size();
        let frame_len = config.frame_length_samples();

        // Povey window: like hamming but goes to zero at edges
        // Formula: pow(0.5 - 0.5*cos(2*PI*i/(N-1)), 0.85)
        let window: Vec<f64> = (0..frame_len)
            .map(|i| {
                let a = 2.0 * PI * i as f64 / (frame_len - 1) as f64;
                (0.5 - 0.5 * a.cos()).powf(0.85)
            })
            .collect();

        // Mel filterbank
        let high_freq = if config.high_freq == 0.0 {
            config.sample_rate / 2.0
        } else {
            config.high_freq
        };
        let mel_filters = kaldi_mel_filterbank(
            config.sample_rate,
            fft_size,
            config.num_mel_bins,
            config.low_freq,
            high_freq,
        );
        let sparse_mel_filters = SparseMelFilterbank::from_dense(&mel_filters);

        let mut planner = RealFftPlanner::new();
        let fft = planner.plan_fft_forward(fft_size);

        Self {
            config,
            mel_filters,
            sparse_mel_filters,
            fft,
            window,
        }
    }

    /// Extract filterbank features from audio samples.
    ///
    /// # Arguments
    /// * `samples` - Audio samples (mono, f32, at the configured sample rate)
    ///
    /// # Returns
    /// * `Array2<f32>` - Filterbank features with shape (num_frames, num_mel_bins)
    pub fn compute(&self, samples: &[f32]) -> Array2<f32> {
        let frame_len = self.config.frame_length_samples();
        let frame_shift = self.config.frame_shift_samples();
        let num_mel_bins = self.config.num_mel_bins;
        let preemph = self.config.preemphasis;

        if samples.len() < frame_len {
            return Array2::zeros((0, num_mel_bins));
        }

        let num_frames = 1 + (samples.len() - frame_len) / frame_shift;
        if num_mel_bins == 0 {
            return Array2::zeros((num_frames, 0));
        }
        let mut features = vec![0.0f32; num_frames * num_mel_bins];

        let mut fft_input = self.fft.make_input_vec();
        let mut spectrum = self.fft.make_output_vec();
        let mut scratch_buf = self.fft.make_scratch_vec();
        let mut frame_buf = vec![0.0f64; frame_len];
        // Power spectra and mel energies of FRAME_LANES frames, interleaved
        // by bin and by mel, for the batched projection.
        let mut power_spectra = vec![0.0f64; spectrum.len() * FRAME_LANES];
        let mut mel_energies = vec![0.0f64; num_mel_bins * FRAME_LANES];

        // Kaldi uses FLT_EPSILON (f32::EPSILON ≈ 1.19e-7) as minimum to avoid log(0)
        let floor = if self.config.energy_floor > 0.0 {
            self.config.energy_floor
        } else {
            f32::EPSILON as f64 // ~1.19e-7, matches kaldi FLT_EPSILON
        };

        for (group_idx, group_features) in
            features.chunks_mut(num_mel_bins * FRAME_LANES).enumerate()
        {
            for lane in 0..group_features.len() / num_mel_bins {
                let start = ((group_idx * FRAME_LANES) + lane) * frame_shift;
                let end = start + frame_len;

                // Copy frame and subtract mean (DC removal)
                let frame_slice = &samples[start..end];
                let mean: f64 =
                    frame_slice.iter().map(|&x| x as f64).sum::<f64>() / frame_len as f64;
                for (i, &sample) in frame_slice.iter().enumerate() {
                    frame_buf[i] = sample as f64 - mean;
                }

                // Apply preemphasis: y[n] = x[n] - preemph * x[n-1]
                if preemph > 0.0 {
                    // Process in reverse to avoid overwriting
                    for i in (1..frame_len).rev() {
                        frame_buf[i] -= preemph * frame_buf[i - 1];
                    }
                    // First sample: use sample from before this frame if available
                    if start > 0 {
                        frame_buf[0] -= preemph * (samples[start - 1] as f64 - mean);
                    }
                }

                // Apply window and zero-pad to FFT size
                for ((input, sample), weight) in
                    fft_input.iter_mut().zip(&frame_buf).zip(&self.window)
                {
                    *input = sample * weight;
                }
                fft_input[frame_len..].fill(0.0);

                // FFT (only positive frequencies)
                self.fft
                    .process_with_scratch(&mut fft_input, &mut spectrum, &mut scratch_buf)
                    .expect("buffers are sized for the FFT plan");

                let lanes = power_spectra.chunks_exact_mut(FRAME_LANES).zip(&spectrum);
                if self.config.use_power {
                    for (frames, c) in lanes {
                        frames[lane] = c.norm_sqr();
                    }
                } else {
                    for (frames, c) in lanes {
                        frames[lane] = c.norm();
                    }
                }
            }

            // Lanes past the last frame hold stale spectra and are not read.
            self.sparse_mel_filters
                .project_frames_f64::<FRAME_LANES>(&power_spectra, &mut mel_energies);
            for (lane, frame_features) in group_features.chunks_exact_mut(num_mel_bins).enumerate()
            {
                for (feature, energies) in frame_features
                    .iter_mut()
                    .zip(mel_energies.chunks_exact(FRAME_LANES))
                {
                    // Apply energy floor and log
                    let mut mel_energy = energies[lane].max(floor);
                    if self.config.use_log_fbank {
                        mel_energy = mel_energy.ln();
                    }
                    *feature = mel_energy as f32;
                }
            }
        }

        // CMN: subtract mean across time for each frequency bin
        // This matches kaldi's CMN: features - np.mean(features, axis=0)
        // The sums run row by row, which keeps the frame order of each
        // column sum and reads the features contiguously.
        if self.config.apply_cmn && num_frames > 0 {
            let mut means = vec![0.0f32; num_mel_bins];
            if num_mel_bins == 1 {
                // A single column is contiguous, and ndarray sums contiguous
                // data with an unrolled fold. This keeps that summation order.
                means[0] = ArrayView1::from(&features[..]).mean().unwrap_or(0.0);
            } else {
                for frame_features in features.chunks_exact(num_mel_bins) {
                    for (sum, feature) in means.iter_mut().zip(frame_features) {
                        *sum += feature;
                    }
                }
                for mean in means.iter_mut() {
                    *mean /= num_frames as f32;
                }
            }
            for frame_features in features.chunks_exact_mut(num_mel_bins) {
                for (feature, mean) in frame_features.iter_mut().zip(&means) {
                    *feature -= mean;
                }
            }
        }

        Array2::from_shape_vec((num_frames, num_mel_bins), features)
            .expect("feature buffer matches the frame and mel counts")
    }

    /// Get the configuration.
    pub fn config(&self) -> &FbankConfig {
        &self.config
    }

    /// Dense Kaldi-style filterbank weights used as the reference projection.
    pub fn dense_filterbank(&self) -> &Array2<f64> {
        &self.mel_filters
    }
}

/// Create Kaldi-style mel filterbank.
///
/// Uses Kaldi mel scale: mel = 1127 * ln(1 + hz/700)
/// Filters are NOT area-normalized (matching kaldi default).
fn kaldi_mel_filterbank(
    sample_rate: f64,
    fft_size: usize,
    num_mel_bins: usize,
    low_freq: f64,
    high_freq: f64,
) -> Array2<f64> {
    let num_fft_bins = fft_size / 2 + 1;

    let mel_low = hz_to_mel(low_freq);
    let mel_high = hz_to_mel(high_freq);

    // Mel bin edges (num_mel_bins + 2 for triangular filters)
    let mel_points: Vec<f64> = (0..=num_mel_bins + 1)
        .map(|i| mel_low + (mel_high - mel_low) * i as f64 / (num_mel_bins + 1) as f64)
        .collect();

    // Convert mel points back to Hz
    let hz_points: Vec<f64> = mel_points.iter().map(|&m| mel_to_hz(m)).collect();

    // Build triangular filters using continuous frequency values
    // Each FFT bin corresponds to a frequency: bin_freq = bin_idx * sample_rate / fft_size
    let mut filters = Array2::zeros((num_mel_bins, num_fft_bins));

    for mel_idx in 0..num_mel_bins {
        let left_hz = hz_points[mel_idx];
        let center_hz = hz_points[mel_idx + 1];
        let right_hz = hz_points[mel_idx + 2];

        // Skip degenerate filters
        if center_hz <= left_hz || right_hz <= center_hz {
            continue;
        }

        for freq_idx in 0..num_fft_bins {
            let freq_hz = freq_idx as f64 * sample_rate / fft_size as f64;

            if freq_hz > left_hz && freq_hz <= center_hz {
                // Rising edge
                filters[[mel_idx, freq_idx]] = (freq_hz - left_hz) / (center_hz - left_hz);
            } else if freq_hz > center_hz && freq_hz < right_hz {
                // Falling edge
                filters[[mel_idx, freq_idx]] = (right_hz - freq_hz) / (right_hz - center_hz);
            }
        }
    }

    filters
}

/// Convert Hz to Mel scale (Kaldi formula).
/// mel = 1127 * ln(1 + hz/700)
fn hz_to_mel(hz: f64) -> f64 {
    1127.0 * (1.0 + hz / 700.0).ln()
}

/// Convert Mel to Hz scale (Kaldi formula).
/// hz = 700 * (exp(mel/1127) - 1)
fn mel_to_hz(mel: f64) -> f64 {
    700.0 * ((mel / 1127.0).exp() - 1.0)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_support::{fuzz_cases, same_f32, Rng};
    use ndarray_npy::NpzReader;
    use std::fs::File;
    use std::io::Read;

    /// Find the start of audio data in a WAV file by locating the "data" chunk.
    /// Returns the byte offset where sample data begins.
    fn find_wav_data_offset(wav_bytes: &[u8]) -> Option<usize> {
        // WAV files have RIFF header (12 bytes) followed by chunks
        // Each chunk: 4-byte ID + 4-byte size + data
        if wav_bytes.len() < 12 {
            return None;
        }

        let mut pos = 12; // Skip RIFF header
        while pos + 8 <= wav_bytes.len() {
            let chunk_id = &wav_bytes[pos..pos + 4];
            let chunk_size = u32::from_le_bytes([
                wav_bytes[pos + 4],
                wav_bytes[pos + 5],
                wav_bytes[pos + 6],
                wav_bytes[pos + 7],
            ]) as usize;

            if chunk_id == b"data" {
                return Some(pos + 8); // Data starts after chunk header
            }

            pos += 8 + chunk_size;
            // Chunks are word-aligned (2-byte boundary)
            if chunk_size % 2 == 1 {
                pos += 1;
            }
        }
        None
    }

    #[test]
    fn test_fbank_config_defaults() {
        let config = FbankConfig::default();
        assert_eq!(config.sample_rate, 16000.0);
        assert_eq!(config.num_mel_bins, 80);
        assert_eq!(config.frame_length_samples(), 400);
        assert_eq!(config.frame_shift_samples(), 160);
        assert_eq!(config.fft_size(), 512);
    }

    #[test]
    fn test_hz_to_mel() {
        // Test values: kaldi uses mel = 1127 * ln(1 + hz/700)
        assert!((hz_to_mel(0.0) - 0.0).abs() < 1e-6);
        // 1000 Hz -> 1127 * ln(1 + 1000/700) = 1127 * ln(2.4286) = 999.98
        assert!((hz_to_mel(1000.0) - 999.98).abs() < 1.0);
        // 8000 Hz -> 1127 * ln(1 + 8000/700) = 1127 * ln(12.4286) = 2840.02
        assert!((hz_to_mel(8000.0) - 2840.0).abs() < 1.0);
    }

    #[test]
    fn test_mel_to_hz() {
        // Round-trip test
        for hz in [0.0, 500.0, 1000.0, 4000.0, 8000.0] {
            let mel = hz_to_mel(hz);
            let hz_back = mel_to_hz(mel);
            assert!(
                (hz - hz_back).abs() < 1e-6,
                "Round-trip failed for Hz={}",
                hz
            );
        }
    }

    #[test]
    fn test_fbank_basic() {
        let config = FbankConfig::default();
        let fbank = Fbank::new(config);

        // Create a simple test signal (1 second of silence)
        let samples = vec![0.0f32; 16000];
        let features = fbank.compute(&samples);

        // Check output shape
        // For 1 second at 16kHz with 25ms frames and 10ms shift:
        // num_frames = 1 + (16000 - 400) / 160 = 98
        assert_eq!(features.shape()[1], 80); // num_mel_bins
        assert!(features.shape()[0] > 90 && features.shape()[0] < 100);
    }

    #[test]
    fn test_sparse_projection_matches_dense_filterbank() {
        let config = FbankConfig::default();
        let fbank = Fbank::new(config);
        let power_spectrum = (0..fbank.mel_filters.ncols())
            .map(|idx| ((idx as f64 + 1.0) * 0.013).sin().abs())
            .collect::<Vec<_>>();
        let mut sparse = vec![0.0; fbank.config.num_mel_bins];

        fbank
            .sparse_mel_filters
            .project_power_f64(&power_spectrum, &mut sparse);

        for (mel_idx, sparse) in sparse.iter().enumerate() {
            let dense = fbank
                .mel_filters
                .row(mel_idx)
                .iter()
                .zip(power_spectrum.iter())
                .map(|(filter, power)| filter * power)
                .sum::<f64>();
            assert!(
                (sparse - dense).abs() <= 1e-12,
                "mel {mel_idx}: sparse {sparse}, dense {dense}"
            );
        }

        assert!(
            fbank.sparse_mel_filters.non_zero_weights()
                < fbank.sparse_mel_filters.dense_weights() / 10
        );
    }

    /// Frame-by-frame form of `Fbank::compute` with the dense filterbank.
    fn reference_fbank(fbank: &Fbank, samples: &[f32]) -> Array2<f32> {
        let config = fbank.config();
        let frame_len = config.frame_length_samples();
        let frame_shift = config.frame_shift_samples();
        let num_frames = 1 + (samples.len() - frame_len) / frame_shift;
        let dense = fbank.dense_filterbank();
        let fft = RealFftPlanner::<f64>::new().plan_fft_forward(config.fft_size());
        let floor = if config.energy_floor > 0.0 {
            config.energy_floor
        } else {
            f32::EPSILON as f64
        };
        let mut features = Array2::<f32>::zeros((num_frames, config.num_mel_bins));
        for frame_idx in 0..num_frames {
            let start = frame_idx * frame_shift;
            let frame = &samples[start..start + frame_len];
            let mean = frame.iter().map(|&x| x as f64).sum::<f64>() / frame_len as f64;
            let mut buf = frame.iter().map(|&x| x as f64 - mean).collect::<Vec<_>>();
            if config.preemphasis > 0.0 {
                for i in (1..frame_len).rev() {
                    buf[i] -= config.preemphasis * buf[i - 1];
                }
                if start > 0 {
                    buf[0] -= config.preemphasis * (samples[start - 1] as f64 - mean);
                }
            }
            let mut input = fft.make_input_vec();
            for i in 0..frame_len {
                input[i] = buf[i] * fbank.window[i];
            }
            let mut spectrum = fft.make_output_vec();
            fft.process(&mut input, &mut spectrum).unwrap();
            for mel_idx in 0..config.num_mel_bins {
                let mut energy = 0.0;
                for (weight, c) in dense.row(mel_idx).iter().zip(&spectrum) {
                    let power = if config.use_power {
                        c.norm_sqr()
                    } else {
                        c.norm()
                    };
                    energy += weight * power;
                }
                let mut energy = f64::max(energy, floor);
                if config.use_log_fbank {
                    energy = energy.ln();
                }
                features[[frame_idx, mel_idx]] = energy as f32;
            }
        }
        if config.apply_cmn {
            for mel_idx in 0..config.num_mel_bins {
                let mean = features.column(mel_idx).mean().unwrap();
                features
                    .column_mut(mel_idx)
                    .mapv_inplace(|value| value - mean);
            }
        }
        features
    }

    #[test]
    fn fbank_matches_frame_by_frame_dense_reference() {
        let samples = (0..16_000)
            .map(|idx| {
                ((idx as f32) * 0.05).sin() * 0.3 + ((idx * 7919 % 997) as f32 / 997.0 - 0.5) * 0.05
            })
            .collect::<Vec<_>>();
        let configs = [
            FbankConfig::default(),
            FbankConfig {
                apply_cmn: false,
                use_power: false,
                energy_floor: 1e-3,
                ..FbankConfig::default()
            },
            FbankConfig {
                num_mel_bins: 23,
                use_log_fbank: false,
                preemphasis: 0.0,
                ..FbankConfig::default()
            },
        ];
        for config in configs {
            let fbank = Fbank::new(config);
            for len in [400, 401, 560, 1_840, 16_000] {
                let got = fbank.compute(&samples[..len]);
                let want = reference_fbank(&fbank, &samples[..len]);
                assert_eq!(got, want, "len {len}");
            }
        }
    }

    #[test]
    fn fuzz_fbank_matches_frame_by_frame_dense_reference() {
        let mut rng = Rng::new(7);
        for case in 0..fuzz_cases(60) {
            let sample_rate = [8_000.0, 16_000.0, 22_050.0, 44_100.0][rng.range(0, 3)];
            let min_ms = 2_000.0 / sample_rate;
            let low_freq = if rng.chance(0.5) {
                20.0
            } else {
                rng.unit() * 200.0
            };
            let config = FbankConfig {
                sample_rate,
                num_mel_bins: if rng.chance(0.2) { 1 } else { rng.range(1, 64) },
                frame_length_ms: min_ms + (rng.unit() * 40.0),
                frame_shift_ms: min_ms + (rng.unit() * 25.0),
                dither: 0.0,
                energy_floor: if rng.chance(0.5) { 0.0 } else { rng.unit() },
                use_energy: false,
                use_log_fbank: rng.chance(0.8),
                use_power: rng.chance(0.8),
                preemphasis: [0.0, 0.97, rng.unit()][rng.range(0, 2)],
                apply_cmn: rng.chance(0.6),
                low_freq,
                high_freq: if rng.chance(0.5) {
                    0.0
                } else {
                    low_freq + (rng.unit() * (sample_rate / 2.0 - low_freq))
                },
            };
            let fbank = Fbank::new(config.clone());
            let frame_len = config.frame_length_samples();
            let len = frame_len + rng.range(0, frame_len * 12);
            let samples = rng.signal(len);
            let got = fbank.compute(&samples);
            let want = reference_fbank(&fbank, &samples);
            assert_eq!(got.dim(), want.dim(), "case {case}: {config:?}");
            for (idx, (got, want)) in got.iter().zip(want.iter()).enumerate() {
                assert!(
                    same_f32(*got, *want),
                    "case {case} value {idx}: {got} vs {want}, {config:?} len {len}"
                );
            }
        }
    }

    #[test]
    fn test_fbank_vs_kaldi_golden() {
        // Load golden data from kaldi_native_fbank
        // Note: This test is informational - our implementation is an approximation
        // and may not exactly match kaldi_native_fbank. For exact compatibility,
        // use a TorchScript-traced version of torchaudio.compliance.kaldi.fbank.
        let npz_path = "./testdata/kaldi_native_fbank_jfk.npz";
        if !std::path::Path::new(npz_path).exists() {
            eprintln!("Skipping golden test: {} not found", npz_path);
            return;
        }

        let f = File::open(npz_path).unwrap();
        let mut npz = NpzReader::new(f).unwrap();
        let golden: Array2<f32> = npz.by_name("features").unwrap();

        // Load the audio file
        let wav_path = "./testdata/jfk_f32le.wav";
        let mut wav_file = File::open(wav_path).unwrap();
        let mut wav_bytes = Vec::new();
        wav_file.read_to_end(&mut wav_bytes).unwrap();

        // Find the "data" chunk in WAV file (handles extended headers)
        let data_offset =
            find_wav_data_offset(&wav_bytes).expect("Could not find 'data' chunk in WAV file");
        let samples: Vec<f32> = wav_bytes[data_offset..]
            .chunks_exact(4)
            .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
            .collect();

        // Compute fbank features with CMN (matching kaldi.py)
        let config = FbankConfig {
            apply_cmn: true,
            ..FbankConfig::default()
        };
        let fbank = Fbank::new(config);
        let computed = fbank.compute(&samples);

        // Golden data is transposed (80, num_frames) -> need (num_frames, 80)
        let golden_t = golden.t();

        eprintln!("Computed shape: {:?}", computed.shape());
        eprintln!("Golden shape: {:?}", golden_t.shape());

        // Check shape matches - this IS required
        assert_eq!(
            computed.shape()[0],
            golden_t.shape()[0],
            "Frame count mismatch: computed {} vs golden {}",
            computed.shape()[0],
            golden_t.shape()[0]
        );

        // Compute differences (informational)
        let num_check = computed.shape()[0].min(50);
        let mut max_diff = 0.0f32;
        let mut sum_diff = 0.0f32;
        let mut count = 0;

        for frame_idx in 0..num_check {
            for mel_idx in 0..80 {
                let diff = (computed[[frame_idx, mel_idx]] - golden_t[[frame_idx, mel_idx]]).abs();
                max_diff = max_diff.max(diff);
                sum_diff += diff;
                count += 1;
            }
        }

        let avg_diff = sum_diff / count as f32;
        eprintln!("Max difference: {:.4}", max_diff);
        eprintln!("Avg difference: {:.4}", avg_diff);

        // Log first frame comparison for debugging
        eprintln!("\nFirst frame comparison (computed vs golden):");
        for mel_idx in 0..5 {
            eprintln!(
                "  mel[{}]: {:.4} vs {:.4}",
                mel_idx,
                computed[[0, mel_idx]],
                golden_t[[0, mel_idx]]
            );
        }

        // NOTE: This implementation differs from kaldi_native_fbank.
        // The test passes as long as we produce valid output with correct shape.
        // For exact kaldi compatibility, use TorchScript-traced fbank.
        eprintln!("\nNote: This is an approximation of kaldi fbank.");
        eprintln!("For exact kaldi compatibility, use TorchScript-traced fbank model.");

        // Verify we produce finite, reasonable values
        let all_finite = computed.iter().all(|&x| x.is_finite());
        assert!(all_finite, "Computed features contain non-finite values");

        // Verify some variation in output (not all zeros or constant)
        let variance: f32 = computed.iter().map(|&x| x * x).sum::<f32>() / computed.len() as f32;
        assert!(variance > 0.1, "Output variance too low: {}", variance);
    }

    #[test]
    fn debug_filterbank() {
        let config = FbankConfig::default();
        let fbank = Fbank::new(config);

        // Check filterbank weights
        println!("\nFilterbank check:");
        for mel_idx in 0..10 {
            let row = fbank.mel_filters.row(mel_idx);
            let sum: f64 = row.iter().sum();
            let nonzero: usize = row.iter().filter(|&&x| x > 0.0).count();
            println!(
                "  Filter {}: sum={:.4}, nonzero_bins={}",
                mel_idx, sum, nonzero
            );
        }

        // Check first few filter shapes
        println!("\nFirst 3 filters (non-zero weights):");
        for mel_idx in 0..3 {
            let row = fbank.mel_filters.row(mel_idx);
            print!("  Filter {}: ", mel_idx);
            for (i, &w) in row.iter().enumerate() {
                if w > 0.0 {
                    print!("bin{}={:.3} ", i, w);
                }
            }
            println!();
        }
    }

    #[test]
    fn debug_compute_steps() {
        // Load audio
        let mut wav_file = File::open("./testdata/jfk_f32le.wav").unwrap();
        let mut wav_bytes = Vec::new();
        wav_file.read_to_end(&mut wav_bytes).unwrap();
        let data_offset =
            find_wav_data_offset(&wav_bytes).expect("Could not find 'data' chunk in WAV file");
        let samples: Vec<f32> = wav_bytes[data_offset..]
            .chunks_exact(4)
            .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
            .collect();

        println!("\nFirst 10 audio samples: {:?}", &samples[..10]);

        // Compute without CMN
        let config = FbankConfig {
            apply_cmn: false,
            ..FbankConfig::default()
        };
        let fbank = Fbank::new(config);
        let features = fbank.compute(&samples);

        println!("\nFrame 0 (silent), first 5 mel bins (no CMN):");
        for i in 0..5 {
            println!("  mel[{}]: {:.4}", i, features[[0, i]]);
        }

        println!("\nExpected (from kaldi): -15.94 for all");
        println!(
            "f32::EPSILON: {:e}, ln(EPSILON): {:.4}",
            f32::EPSILON,
            (f32::EPSILON as f64).ln()
        );
    }
}