lasprs 0.14.2

//! Frequency-domain smoothing of power spectra.
//!
//! Implements fractional-octave smoothing on single-sided power spectra in the
//! frequency domain. The algorithm interpolates the spectrum to a logarithmic
//! frequency grid, applies a symmetric moving-average window whose width
//! corresponds to the requested fraction of an octave, and back-interpolates
//! the result onto the original linear frequency grid.
//!
//! The main entry point for users is [`smoothSpectralData`], which handles
//! dB ↔ power conversion based on the [`SmoothingType`]. The lower-level
//! [`freqSmooth`] operates on power data directly.

use super::error::*;
use crate::config::*;
use ndarray::par_azip;
use ndarray_interp::interp1d::{Interp1DBuilder, Linear};
use rayon::prelude::*;
use snafu::prelude::*;
#[cfg(feature = "python-bindings")]
use strum::EnumMessage;
use strum::IntoEnumIterator;
use strum_macros::{Display, EnumIter, EnumMessage};

// ── SmoothingType enum ────────────────────────────────────────────────────

/// The type of spectral data being smoothed. Determines how the input is
/// interpreted and what conversions are applied before / after the actual
/// frequency-domain smoothing.
#[derive(Copy, Clone, Debug, Default, PartialEq, Hash, Display, EnumIter, EnumMessage)]
#[cfg_attr(
    feature = "python-bindings",
    gen_stub_pyclass_enum,
    pyclass(eq, eq_int, from_py_object)
)]
pub enum SmoothingType {
    /// Input data is in decibels (levels). It is converted to power
    /// (10^(M/10)) before smoothing, and converted back to dB afterwards.
    #[default]
    #[strum(message = "Levels (dB)")]
    Levels = 0,
    /// Input data is already in (auto) power units. Smoothing is applied
    /// directly.
    #[strum(message = "(Auto) powers")]
    Power = 1,
}

#[cfg(feature = "python-bindings")]
#[cfg_attr(feature = "python-bindings", gen_stub_pymethods, pymethods)]
impl SmoothingType {
    #[staticmethod]
    fn all() -> Vec<SmoothingType> {
        SmoothingType::iter().collect()
    }
    fn __str__(&self) -> String {
        format!("{self}")
    }
    #[staticmethod]
    #[pyo3(name = "default")]
    fn default_py() -> Self {
        Self::default()
    }
}
// ── SmoothingWidth enum ────────────────────────────────────────────────────

/// Fractional-octave smoothing width. Determines how wide the moving-average
/// window is when smoothing a power spectrum on a logarithmic frequency grid.
///
/// `Oct1` corresponds to full-octave (1/1) smoothing, `Oct3` to 1/3-octave,
/// and so on up to `Oct48` for 1/48-octave. `Oct100` is provided as a very
/// narrow smoothing option. [`None`](SmoothingWidth::None) disables smoothing
/// entirely and returns the input unchanged.
#[derive(Copy, Clone, Debug, Default, PartialEq, Hash, Display, EnumIter, EnumMessage)]
#[cfg_attr(
    feature = "python-bindings",
    gen_stub_pyclass_enum,
    pyclass(eq, eq_int, from_py_object)
)]
pub enum SmoothingWidth {
    /// No smoothing — input is returned as-is.
    #[strum(message = "No smoothing")]
    NoSmoothing = 0,
    /// 1/1-octave smoothing
    #[strum(message = "1/1 octave")]
    Oct1 = 1,
    /// 1/2-octave smoothing
    #[strum(message = "1/2 octave")]
    Oct2 = 2,
    /// 1/3-octave smoothing
    #[default]
    #[strum(message = "1/3 octave")]
    Oct3 = 3,
    /// 1/4-octave smoothing
    #[strum(message = "1/4 octave")]
    Oct4 = 4,
    /// 1/6-octave smoothing
    #[strum(message = "1/6 octave")]
    Oct6 = 6,
    /// 1/8-octave smoothing
    #[strum(message = "1/8 octave")]
    Oct8 = 8,
    /// 1/12-octave smoothing
    #[strum(message = "1/12 octave")]
    Oct12 = 12,
    /// 1/16-octave smoothing
    #[strum(message = "1/16 octave")]
    Oct16 = 16,
    /// 1/24-octave smoothing
    #[strum(message = "1/24 octave")]
    Oct24 = 24,
    /// 1/48-octave smoothing
    #[strum(message = "1/48 octave")]
    Oct48 = 48,
    /// 1/100-octave smoothing (very narrow)
    #[strum(message = "1/100 octave")]
    Oct100 = 100,
}

impl SmoothingWidth {
    /// Return the denominator of the fractional-octave width, or `0` for
    /// [`SmoothingWidth::None`].
    ///
    /// For example, `Oct3` returns `3`, meaning 1/3-octave.
    pub fn w(&self) -> usize {
        *self as usize
    }

    /// Returns `true` when this width actually requests smoothing.
    pub fn is_smoothing(&self) -> bool {
        *self != SmoothingWidth::NoSmoothing
    }
}

#[cfg(feature = "python-bindings")]
#[cfg_attr(feature = "python-bindings", gen_stub_pymethods, pymethods)]
impl SmoothingWidth {
    #[staticmethod]
    fn all() -> Vec<SmoothingWidth> {
        SmoothingWidth::iter().collect()
    }
    fn __str__(&self) -> String {
        self.get_message().unwrap().into()
    }
    #[staticmethod]
    #[pyo3(name = "default")]
    fn default_py() -> Self {
        Self::default()
    }
}

// ── Private helpers ────────────────────────────────────────────────────────

type Result<T> = std::result::Result<T, FreqSmoothError>;

/// Build an `ndarray-interp` linear interpolator (with extrapolation enabled)
/// from known `(xp, yp)` data and interpolate at the query points `xq`.
fn interp_linear(xp: ArrayView1<Flt>, yp: ArrayView1<Flt>, xq: ArrayView1<Flt>) -> Result<Dcol> {
    let interpolator = Interp1DBuilder::new(yp)
        .x(xp)
        .strategy(Linear::new().extrapolate(true))
        .build()
        .map_err(|e| InterpolationFailedSnafu { msg: e.to_string() }.build())?;

    interpolator
        .interp_array(&xq)
        .map_err(|e| InterpolationFailedSnafu { msg: e.to_string() }.build())
}

// ── freqSmooth (low-level, operates on power data) ─────────────────────────

/// Apply fractional-octave frequency-domain smoothing to a single-sided power
/// spectrum.
///
/// This is the low-level smoothing routine. Most callers should prefer
/// [`smoothSpectralData`], which additionally handles dB ↔ power conversion.
///
/// # Arguments
///
/// * `freq`  – Frequency vector (Hz). Must be monotonically increasing. The
///   first element may be 0 (DC).
/// * `X`     – Power spectrum values corresponding to each frequency bin.
/// * `smoothing_width` – Fractional-octave smoothing width (see
///   [`SmoothingWidth`]). Must not be [`SmoothingWidth::None`] — callers
///   should check [`SmoothingWidth::is_smoothing`] before calling.
/// * `power_correct` – When `true`, the smoothed spectrum is rescaled so that
///   the total AC power equals that of the original spectrum.
///
/// # Returns
///
/// A new [`Dcol`] with the smoothed power spectrum, same length as `freq`.
///
/// # Errors
///
/// Returns a [`FreqSmoothError`] when any of the input constraints are
/// violated (sizes, frequency range, etc.) or if interpolation fails.
fn freqSmooth(
    freq: ArrayView1<Flt>,
    X: ArrayView1<Flt>,
    smoothing_width: SmoothingWidth,
    power_correct: bool,
) -> Result<Dcol> {
    let Nfreq = freq.len();
    let w = smoothing_width.w();

    // ── Handle DC bin ──────────────────────────────────────────────────
    let firstFreqEqZero = freq[0].abs() < 1e-15;

    // Lowest nonzero frequency
    let freq_min: Flt;
    let freq_max: Flt = freq[Nfreq - 1];

    // Optionally compute AC power before smoothing
    let ac_pwr: Flt;
    if firstFreqEqZero {
        freq_min = freq[1];
        ac_pwr = if power_correct {
            X.slice(ndarray::s![1..]).sum()
        } else {
            0.
        };
    } else {
        freq_min = freq[0];
        ac_pwr = if power_correct { X.sum() } else { 0. };
    }

    ensure!(freq_min > 0., InvalidFreqMinSnafu { freq_min });

    // ── Build logarithmic frequency grid (10× oversampling) ────────────
    let Nfreq_sm = 10 * Nfreq;
    let log_freq_min = freq_min.log10();
    let log_freq_max = freq_max.log10();
    let freq_log = Dcol::from_iter((0..Nfreq_sm).map(|i| {
        let t = i as Flt / (Nfreq_sm - 1) as Flt;
        Flt::powf(10., log_freq_min + t * (log_freq_max - log_freq_min))
    }));

    // ── Forward interpolation: linear freq → log freq ──────────────────
    let mut X_log = interp_linear(freq, X, freq_log.view())?;

    // Fix boundary points that may be slightly out of range due to
    // floating-point round-off.
    X_log[Nfreq_sm - 1] = X[X.len() - 1];
    if firstFreqEqZero {
        X_log[0] = X[1];
    } else {
        X_log[0] = X[0];
    }

    // ── Compute smoothing half-width in index units ────────────────────
    let Delta: Flt = 1. / w as Flt; // smoothing width in octaves
    let fstep: Flt = freq_log[1] / freq_log[0]; // ratio between consecutive log bins
    let hicenter: Flt = Flt::powf(2., Delta / 2.); // f_high / f_center
    let mu: usize = (hicenter.log10() / fstep.log10()) as usize;

    // ── Symmetric moving-average on log grid ───────────────────────────
    let mut Xsm_log = Dcol::zeros(Nfreq_sm);
    Xsm_log
        .as_slice_mut()
        .expect("Cannot slice?")
        // Hurray for Rayon. This one line makes the code run this smoothing in
        // parallel.
        .par_iter_mut()
        .enumerate()
        .for_each(|(k, Xsm_log)| {
            let mut idx_start = k.saturating_sub(mu);
            let mut idx_stop = (k + mu).min(Nfreq_sm - 1);

            // Shrink window symmetrically at edges
            if idx_start == 0 || idx_stop == Nfreq_sm - 1 {
                let mu_edge = (k - idx_start).min(idx_stop - k);
                idx_start = k - mu_edge;
                idx_stop = k + mu_edge;
            }
            let slice = X_log.slice(ndarray::s![idx_start..=idx_stop]);
            *Xsm_log = slice.mean().unwrap_or(X_log[k]);
        });

    // ── Back-interpolation: log freq → linear freq ─────────────────────
    let mut Xsm = Dcol::zeros(Nfreq);

    if firstFreqEqZero {
        // Interpolate only the non-DC portion
        let freq_gt0 = freq.slice(ndarray::s![1..]).to_owned();
        let Xsm_gt0 = interp_linear(freq_log.view(), Xsm_log.view(), freq_gt0.view())?;

        Xsm[0] = X[0]; // preserve DC bin
        Xsm.slice_mut(ndarray::s![1..]).assign(&Xsm_gt0.view());

        // Fix boundary points
        Xsm[1] = Xsm_log[1];
        Xsm[Nfreq - 1] = Xsm_log[Nfreq_sm - 1];

        // Power correction
        if power_correct {
            let new_acpwr: Flt = Xsm.slice(ndarray::s![1..]).sum();
            if new_acpwr.abs() > 1e-30 {
                let scale = ac_pwr / new_acpwr;
                Xsm.slice_mut(ndarray::s![1..]).mapv_inplace(|v| v * scale);
            }
        }
    } else {
        Xsm = interp_linear(freq_log.view(), Xsm_log.view(), freq.view())?;

        // Fix boundary points
        Xsm[0] = X[0];
        Xsm[Nfreq - 1] = Xsm_log[Nfreq_sm - 1];

        // Power correction
        if power_correct {
            let new_acpwr: Flt = Xsm.sum();
            if new_acpwr.abs() > 1e-30 {
                let scale = ac_pwr / new_acpwr;
                Xsm.mapv_inplace(|v| v * scale);
            }
        }
    }

    Ok(Xsm)
}

// ── smoothSpectralData (high-level) ────────────────────────────────────────

/// Apply fractional-octave smoothing to spectral data in the frequency domain.
///
/// This is the main entry point. Depending on the [`SmoothingType`], the input
/// is interpreted either as levels in dB or as (auto-) power values. When the
/// input is in dB, it is first converted to power, smoothed, and then converted
/// back to dB.
///
/// When the [`SmoothingWidth`] is [`None`](SmoothingWidth::None), the input is
/// returned unchanged (after validation).
///
/// # Arguments
///
/// * `freq` – Frequency vector (Hz). Must be monotonically increasing. The
///   first element may be 0 (DC).
/// * `M` – Spectral data. Either levels (dB) or power, depending on `st`.
/// * `sw` – Fractional-octave smoothing width.
/// * `st` – Type of the spectral data (see [`SmoothingType`]).
///
/// # Returns
///
/// A new [`Dcol`] with the smoothed spectral data, in the same units as the
/// input.
///
/// # Errors
///
/// Returns a [`FreqSmoothError`] when inputs are invalid.
pub fn smoothSpectralData(
    freq: ArrayView1<Flt>,
    M: ArrayView1<Flt>,
    sw: SmoothingWidth,
    st: SmoothingType,
) -> Result<Dcol> {
    // ── Input validation ───────────────────────────────────────────────
    ensure!(freq.len() >= 2, FreqTooShortSnafu { length: freq.len() });
    ensure!(
        freq.len() == M.len(),
        SizeMismatchSnafu {
            freq_len: freq.len(),
            x_len: M.len()
        }
    );

    // Type-specific validation
    if st == SmoothingType::Power {
        let min_val = M.fold(Flt::INFINITY, |acc, &v| acc.min(v));
        ensure!(min_val >= 0., NegativePowerSnafu { min_value: min_val });
    }

    // ── No smoothing requested → return input as-is ────────────────────
    if !sw.is_smoothing() {
        return Ok(M.to_owned());
    }

    // ── Convert to power if necessary ──────────────────────────────────
    let P: Dcol = match st {
        SmoothingType::Levels => M.mapv(|v| Flt::powf(10., v / 10.)),
        SmoothingType::Power => M.to_owned(),
    };

    // ── Apply frequency smoothing on power data ────────────────────────
    let Psm = freqSmooth(freq, P.view(), sw, false)?;

    // ── Convert back to original format ────────────────────────────────
    let result = match st {
        SmoothingType::Levels => Psm.mapv(|v| 10. * v.log10()),
        SmoothingType::Power => Psm,
    };

    Ok(result)
}

// ── Python bindings ────────────────────────────────────────────────────────

#[cfg(feature = "python-bindings")]
use numpy::{PyArrayMethods, PyReadonlyArray1};

#[cfg(feature = "python-bindings")]
#[gen_stub_pyfunction]
#[pyfunction(name = "smoothSpectralData")]
/// Apply fractional-octave smoothing to spectral data in the frequency domain.
///
/// Depending on the smoothing type, the input is interpreted either as levels
/// in dB or as (auto-) power values.
///
/// Args:
///     freq: Frequency vector (Hz), monotonically increasing. First element
///         may be 0 (DC).
///     M: Spectral data — either levels (dB) or power values.
///     sw: Fractional-octave smoothing width.
///     st: Type of the spectral data.
///
/// Returns:
///     Smoothed spectral data as a numpy array, in the same units as the input.
pub(crate) fn smoothSpectralData_py<'py>(
    py: Python<'py>,
    freq: PyReadonlyArray1<'py, Flt>,
    M: PyReadonlyArray1<'py, Flt>,
    sw: SmoothingWidth,
    st: SmoothingType,
) -> PyResult<Bound<'py, PyArray1<Flt>>> {
    let freq = freq.as_array();
    let M_arr = M.as_array();
    let result = smoothSpectralData(freq, M_arr, sw, st)?;
    Ok(result.into_pyarray(py))
}

// ── Tests ──────────────────────────────────────────────────────────────────

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

    /// Helper: generate a frequency vector for `nfft` with given sample rate.
    fn make_freq(fs: Flt, nfft: usize) -> Dcol {
        let K = nfft / 2 + 1;
        let df = fs / nfft as Flt;
        Dcol::from_iter((0..K).map(|i| i as Flt * df))
    }

    // ── freqSmooth tests ───────────────────────────────────────────────

    #[test]
    fn test_flat_spectrum_unchanged() {
        let nfft = 1024;
        let fs: Flt = 48000.;
        let freq = make_freq(fs, nfft);
        let X = Dcol::ones(freq.len());

        let Xsm = freqSmooth(freq.view(), X.view(), SmoothingWidth::Oct3, false).unwrap();
        for i in 0..Xsm.len() {
            assert_relative_eq!(Xsm[i], 1.0, epsilon = 1e-6);
        }
    }

    #[test]
    fn test_power_conservation() {
        let nfft = 512;
        let fs: Flt = 44100.;
        let freq = make_freq(fs, nfft);
        let K = freq.len();

        let mut X = Dcol::ones(K);
        X[K / 4] = 100.;

        let Xsm = freqSmooth(freq.view(), X.view(), SmoothingWidth::Oct3, true).unwrap();

        let ac_pwr_orig: Flt = X.slice(ndarray::s![1..]).sum();
        let ac_pwr_sm: Flt = Xsm.slice(ndarray::s![1..]).sum();

        assert_relative_eq!(ac_pwr_orig, ac_pwr_sm, epsilon = 1e-6);
    }

    #[test]
    fn test_dc_preserved() {
        let nfft = 256;
        let fs: Flt = 16000.;
        let freq = make_freq(fs, nfft);
        let mut X = Dcol::ones(freq.len());
        X[0] = 42.;

        let Xsm = freqSmooth(freq.view(), X.view(), SmoothingWidth::Oct1, false).unwrap();
        assert_relative_eq!(Xsm[0], 42., epsilon = 1e-12);
    }

    #[test]
    fn test_smoothing_reduces_peak() {
        let nfft = 1024;
        let fs: Flt = 48000.;
        let freq = make_freq(fs, nfft);
        let K = freq.len();

        let mut X = Dcol::ones(K);
        X[K / 2] = 1000.;

        let Xsm = freqSmooth(freq.view(), X.view(), SmoothingWidth::Oct3, false).unwrap();
        assert!(Xsm[K / 2] < X[K / 2]);
    }

    #[test]
    fn test_invalid_freq_too_short() {
        let freq = Dcol::from_vec(vec![100.]);
        let X = Dcol::from_vec(vec![1.]);
        let err = smoothSpectralData(
            freq.view(),
            X.view(),
            SmoothingWidth::Oct3,
            SmoothingType::Power,
        )
        .unwrap_err();
        assert!(matches!(err, FreqSmoothError::FreqTooShort { .. }));
    }

    #[test]
    fn test_invalid_size_mismatch() {
        let freq = Dcol::from_vec(vec![0., 100., 200.]);
        let X = Dcol::from_vec(vec![1., 2.]);
        let err = smoothSpectralData(
            freq.view(),
            X.view(),
            SmoothingWidth::Oct3,
            SmoothingType::Power,
        )
        .unwrap_err();
        assert!(matches!(err, FreqSmoothError::SizeMismatch { .. }));
    }

    #[test]
    fn test_no_dc() {
        let freq = Dcol::from_iter((1..=100).map(|i| i as Flt * 10.));
        let X = Dcol::ones(100);

        let Xsm = freqSmooth(freq.view(), X.view(), SmoothingWidth::Oct3, false).unwrap();
        assert_eq!(Xsm.len(), 100);
        for i in 0..Xsm.len() {
            assert_relative_eq!(Xsm[i], 1.0, epsilon = 1e-6);
        }
    }

    #[test]
    fn test_different_smoothing_widths() {
        let nfft = 1024;
        let fs: Flt = 48000.;
        let freq = make_freq(fs, nfft);
        let K = freq.len();

        let mut X = Dcol::ones(K);
        X[K / 3] = 500.;

        let Xsm_3 = freqSmooth(freq.view(), X.view(), SmoothingWidth::Oct3, false).unwrap();
        let Xsm_1 = freqSmooth(freq.view(), X.view(), SmoothingWidth::Oct1, false).unwrap();

        // 1/1-octave smoothing (wider) should reduce the peak more than 1/3-octave
        assert!(Xsm_1[K / 3] < Xsm_3[K / 3]);
    }

    #[test]
    fn test_smoothing_width_w() {
        assert_eq!(SmoothingWidth::NoSmoothing.w(), 0);
        assert_eq!(SmoothingWidth::Oct1.w(), 1);
        assert_eq!(SmoothingWidth::Oct3.w(), 3);
        assert_eq!(SmoothingWidth::Oct12.w(), 12);
        assert_eq!(SmoothingWidth::Oct48.w(), 48);
        assert_eq!(SmoothingWidth::Oct100.w(), 100);
    }

    #[test]
    fn test_all_smoothing_widths_work() {
        use strum::IntoEnumIterator;
        let nfft = 512;
        let fs: Flt = 48000.;
        let freq = make_freq(fs, nfft);
        let K = freq.len();
        let mut X = Dcol::ones(K);
        X[K / 4] = 50.;

        for sw in SmoothingWidth::iter() {
            if !sw.is_smoothing() {
                continue;
            }
            let Xsm = freqSmooth(freq.view(), X.view(), sw, false);
            assert!(Xsm.is_ok(), "freqSmooth failed for {:?}", sw);
            assert_eq!(Xsm.unwrap().len(), K);
        }
    }

    // ── smoothSpectralData tests ───────────────────────────────────────

    #[test]
    fn test_smooth_none_returns_input() {
        let nfft = 256;
        let fs: Flt = 48000.;
        let freq = make_freq(fs, nfft);
        let K = freq.len();
        let mut M = Dcol::zeros(K);
        M[K / 4] = 80.; // 80 dB

        let result = smoothSpectralData(
            freq.view(),
            M.view(),
            SmoothingWidth::NoSmoothing,
            SmoothingType::Levels,
        )
        .unwrap();
        assert_eq!(result, M);
    }

    #[test]
    fn test_smooth_levels_roundtrip() {
        // Flat dB input should remain flat after smoothing.
        let nfft = 512;
        let fs: Flt = 48000.;
        let freq = make_freq(fs, nfft);
        let M = Dcol::from_elem(freq.len(), 60.); // 60 dB flat

        let result = smoothSpectralData(
            freq.view(),
            M.view(),
            SmoothingWidth::Oct3,
            SmoothingType::Levels,
        )
        .unwrap();
        for i in 0..result.len() {
            assert_relative_eq!(result[i], 60., epsilon = 1e-3);
        }
    }

    #[test]
    fn test_smooth_power_direct() {
        let nfft = 512;
        let fs: Flt = 48000.;
        let freq = make_freq(fs, nfft);
        let K = freq.len();
        let mut P = Dcol::ones(K);
        P[K / 4] = 100.;

        let result = smoothSpectralData(
            freq.view(),
            P.view(),
            SmoothingWidth::Oct3,
            SmoothingType::Power,
        )
        .unwrap();
        // Peak should be reduced
        assert!(result[K / 4] < P[K / 4]);
        // Off-peak should be increased
        assert!(result[K / 2] >= P[K / 2]);
    }

    #[test]
    fn test_smooth_levels_reduces_peak() {
        let nfft = 1024;
        let fs: Flt = 48000.;
        let freq = make_freq(fs, nfft);
        let K = freq.len();
        let mut M = Dcol::from_elem(K, 40.); // 40 dB floor
        M[K / 3] = 100.; // 100 dB peak

        let result = smoothSpectralData(
            freq.view(),
            M.view(),
            SmoothingWidth::Oct3,
            SmoothingType::Levels,
        )
        .unwrap();
        assert!(result[K / 3] < M[K / 3]);
    }

    #[test]
    fn test_smooth_power_negative_rejected() {
        let freq = Dcol::from_vec(vec![0., 100., 200., 300.]);
        let P = Dcol::from_vec(vec![1., -0.5, 2., 3.]);

        let err = smoothSpectralData(
            freq.view(),
            P.view(),
            SmoothingWidth::Oct3,
            SmoothingType::Power,
        )
        .unwrap_err();
        assert!(matches!(err, FreqSmoothError::NegativePower { .. }));
    }

    #[test]
    fn test_smooth_none_still_validates() {
        // Even when no smoothing, size mismatch should be caught
        let freq = Dcol::from_vec(vec![0., 100., 200.]);
        let M = Dcol::from_vec(vec![1., 2.]);

        let err = smoothSpectralData(
            freq.view(),
            M.view(),
            SmoothingWidth::NoSmoothing,
            SmoothingType::Levels,
        )
        .unwrap_err();
        assert!(matches!(err, FreqSmoothError::SizeMismatch { .. }));
    }
}