lasprs 0.9.1

use super::*;
use super::{timebuffer::TimeBuffer, CrossPowerSpecra};
use crate::{config::*, TransferFunction, ZPKModel};
use anyhow::{bail, Error, Result};
use freqweighting::FreqWeighting;

/// Averaged power spectra computing engine
/// Used to compute power spectra estimations on
/// long datasets, where nfft << length of data. This way, the variance of a
/// single periodogram is suppressed with increasing number of averages.
///
/// For more information, see the book on numerical recipes.

#[cfg_attr(feature = "python-bindings", pyclass)]
#[derive(Debug)]

pub struct AvPowerSpectra {
    // Power spectra estimator for single block
    ps: PowerSpectra,

    // Settings for computing power spectra, see [ApsSettings]
    settings: ApsSettings,

    // The number of samples to keep in the time buffer when overlapping time
    // blocks
    overlap_keep: usize,

    /// The number of blocks of length [self.nfft()] already used in the average
    N: usize,

    /// Storage for sample data.
    timebuf: TimeBuffer,

    /// Power scaling for of applied frequency weighting. Multiply each power
    /// and cross-power value with these constants to apply the frequency
    /// weighting.
    freqWeighting_pwr: Option<Dcol>,

    // Current estimation of the power spectra
    cur_est: CPSResult,
}
impl AvPowerSpectra {
    /// The FFT Length of estimating (cross)power spectra
    pub fn nfft(&self) -> usize {
        self.ps.nfft()
    }

    /// Resets all state, starting with a clean sleave. After this step, also
    /// the number of channels can be different on the input.
    pub fn reset(&mut self) {
        self.N = 0;
        self.timebuf.reset();
        self.cur_est = CPSResult::zeros((0, 0, 0));
    }
    /// Create new averaged power spectra estimator for weighing over the full
    /// amount of data supplied (no exponential spectra weighting) using
    /// sensible defaults (Hann window, 50% overlap). This is a simpler method
    /// than [AvPowerSpectra.build]. But use with caution, it might panic on
    /// invalid nfft values!
    ///
    /// # Args
    ///
    /// * `fs` - Sampling frequency in \[Hz\]
    /// * `nfft` - FFT Length \[-\]
    ///
    /// # Panics
    ///
    /// - When nfft is not even, or 0.
    /// - When providing invalid sampling frequencies
    ///
    pub fn new_simple_all_averaging(fs: Flt, nfft: usize) -> AvPowerSpectra {
        let mut settings = ApsSettings::reasonableAcousticDefault(fs, ApsMode::default()).unwrap();
        settings.nfft = nfft;
        AvPowerSpectra::new(settings)
    }

    /// Create power spectra estimator which weighs either all data
    /// (`fs_tau=None`), or uses exponential weighting. Other parameters can be
    /// provided, like the overlap (Some(Overlap)). Otherwise: all parameters
    /// have sensible defaults.
    ///
    ///  # Exponential weighting
    ///
    ///  This can be used to `follow` a power spectrum as a function of time.
    ///  New data is weighted more heavily. Note that this way of looking at
    ///  spectra is not 'exact', and therefore should not be used for
    ///  spectrograms.
    ///
    /// # Args
    ///
    /// - `nfft` - The discrete Fourier Transform length used in the estimation.
    /// - `windowtype` - Window Type. The window type to use Hann, etc.
    /// - `overlap` - Amount of overlap in
    /// - `mode` - The mode in which the [AvPowerSpectra] runs. See [ApsMode].
    ///
    pub fn new(settings: ApsSettings) -> AvPowerSpectra {
        let overlap_keep = settings.get_overlap_keep();
        let window = Window::new(settings.windowType, settings.nfft);

        let ps = PowerSpectra::newFromWindow(window);

        let freq = settings.getFreq();
        let freqWeighting_pwr = match settings.freqWeightingType {
            FreqWeighting::Z => None,
            _ => {
                let fw_pwr = ZPKModel::freqWeightingFilter(settings.freqWeightingType)
                    .tf(0., &freq)
                    .mapv(|a| a.abs() * a.abs());
                Some(fw_pwr)
            }
        };

        AvPowerSpectra {
            ps,
            overlap_keep,
            settings,
            N: 0,
            freqWeighting_pwr,
            cur_est: CPSResult::default((0, 0, 0)),
            timebuf: TimeBuffer::new(),
        }
    }
    // Update result for single block
    fn update_singleblock(&mut self, timedata: ArrayView2<Flt>) {
        // Compute updated block of cross-spectral density
        let Cpsnew = {
            let mut Cpsnew = self.ps.compute(timedata);
            let dim = Cpsnew.dim();
            if let Some(fw_pwr) = &self.freqWeighting_pwr {
                // Frequency weighting on power is available, multiply all values with frequency weighting

                // Option 1: azip with indexing
                // azip!((index (i,_,_), cps in &mut Cpsnew) {
                //     *cps = Cflt::new(cps.re * fw_pwr[[i]], cps.im * fw_pwr[[i]]);
                // });

                // Option 2: broadcasting with an unwrap.
                let dview_fw = fw_pwr.slice(s![.., NewAxis, NewAxis]);
                azip!((c in &mut Cpsnew, f in dview_fw.broadcast(dim).expect("BUG: Cannot broadcast")) {
                    // Scale with frequency weighting
                    *c = Cflt::new(c.re * f, c.im * f);
                });
            }

            Cpsnew
        };

        // Initialize to zero
        if self.cur_est.is_empty() {
            assert_eq!(self.N, 0);
            self.cur_est = CPSResult::zeros(Cpsnew.raw_dim().f());
        }

        // Update the number of blocks processed
        self.N += 1;

        // Apply operation based on mode
        match self.settings.mode {
            ApsMode::AllAveraging {} => {
                let Nf = Cflt {
                    re: self.N as Flt,
                    im: 0.,
                };
                self.cur_est = (Nf - 1.) / Nf * &self.cur_est + 1. / Nf * Cpsnew;
            }
            ApsMode::ExponentialWeighting { tau } => {
                debug_assert!(self.N > 0);
                if self.N == 1 {
                    self.cur_est = Cpsnew;
                } else {
                    // A sound level meter specifies a low pass filter with one
                    // real pole at -1/tau, for a time weighting of tau. This
                    // means the analogue transfer function is 1 /
                    // (tau*s+1).

                    // Now we want to approximate this with a digital transfer
                    // function. The step size, or sampling period is:
                    // T = (nfft-overlap_keep)/fs.

                    // Then, when using the matched z-transform (
                    // https://en.wikipedia.org/wiki/Matched_Z-transform_method),
                    // an 1/(s-p) is replaced by 1/(1-exp(p*T)*z^-1).

                    // So the digital transfer function will be:
                    // H[n] ≅ K / (1 - exp(-T/tau) * z^-1).
                    // , where K is a to-be-determined constant, for which we
                    // take the value such that the D.C. gain equals 1. To get
                    // the frequency response at D.C., we have to set z=1, so
                    // to set the D.C. to 1, we have to set K = 1-exp(-T/tau).

                    // Hence:
                    // H[n] ≅ (1- exp(-T/tau)) / (1 - exp(-T/tau) * z^-1).

                    // Or as a finite difference equation:
                    //
                    // y[n] * (1-exp(-T/tau)* z^-1) = (1-exp(-T/tau)) * x[n]

                    // or, finally:
                    // y[n] = alpha * y[n-1] + (1-alpha) * x[n]

                    // where alpha = exp(-T/tau).
                    let T = (self.nfft() - self.overlap_keep) as Flt / self.settings.fs;
                    let alpha = Cflt::ONE * Flt::exp(-T / tau);
                    self.cur_est = alpha * &self.cur_est + (1. - alpha) * Cpsnew;
                }
            }

            ApsMode::Spectrogram {} => {
                self.cur_est = Cpsnew;
            }
        }
    }
    /// Computes average (cross)power spectra, and returns only the most recent
    /// estimate, if any can be given back. Only gives back a result when enough
    /// data is available.
    ///
    /// # Args
    ///
    /// * `timedata``: New available time data. Number of columns is number of
    ///   channels, number of rows is number of frames (samples per channel).
    ///
    /// # Panics
    ///
    /// If timedata.ncols() does not match number of columns in already present
    /// data.
    pub fn compute_last<'a, 'b, T>(&'a mut self, timedata: T) -> Option<&'a CPSResult>
    where
        T: AsArray<'b, Flt, Ix2>,
    {
        // Push new data in the time buffer.
        self.timebuf.push(timedata);

        // Flag to indicate that we have obtained one result for sure.
        let mut computed_single = false;

        // Iterate over all blocks that can come,
        while let Some(timeblock) = self.timebuf.pop(self.nfft(), self.overlap_keep) {
            // Compute cross-power spectra for current time block
            self.update_singleblock(timeblock.view());

            computed_single = true;
        }
        if computed_single {
            Some(&self.cur_est)
        } else {
            None
        }
    }

    /// Computes average (cross)power spectra, and returns all intermediate
    /// estimates that can be calculated. This is useful when plotting spectra
    /// as a function of time, and intermediate results need also be plotted.
    ///
    /// # Args
    ///
    /// * `timedata``: New available time data. Number of columns is number of
    ///   channels, number of rows is number of frames (samples per channel).
    ///
    /// # Panics
    ///
    /// If timedata.ncols() does not match number of columns in already present
    /// data.
    pub fn compute_all<'a, 'b, T>(&'a mut self, timedata: T) -> Vec<CPSResult>
    where
        T: AsArray<'b, Flt, Ix2>,
    {
        // Push new data in the time buffer.
        self.timebuf.push(timedata);

        // Storage for the result
        let mut result = Vec::new();

        // Iterate over all blocks that can come,
        while let Some(timeblock) = self.timebuf.pop(self.nfft(), self.overlap_keep) {
            // Compute cross-power spectra for current time block
            self.update_singleblock(timeblock.view());

            result.push(self.cur_est.clone());
        }

        result
    }
}

#[cfg(feature = "python-bindings")]
#[cfg_attr(feature = "python-bindings", pymethods)]
impl AvPowerSpectra {
    #[new]
    fn new_py(s: ApsSettings) -> AvPowerSpectra {
        AvPowerSpectra::new(s)
    }

    #[pyo3(name = "compute")]
    fn compute_py<'py>(
        &mut self,
        py: Python<'py>,
        dat: PyArrayLike2<Flt>,
    ) -> Bound<'py, PyArray3<Cflt>> {
        let dat = dat.as_array();
        if let Some(res) = self.compute_last(dat) {
            let res = res.clone();
            return res.to_pyarray(py);
        }
        let res: Bound<'py, PyArray3<Cflt>> = PyArray3::zeros(py, [0, 0, 0], true);
        res
    }
}
#[cfg(test)]
mod test {
    use approx::assert_abs_diff_eq;

    use super::*;
    use crate::{config::*, math::randNormal};

    use super::{ApsMode, AvPowerSpectra, CPSResult, Overlap, WindowType};
    use Overlap::Percentage;

    #[test]
    fn test_overlap_keep() {
        let ol = [
            Overlap::NoOverlap {},
            Percentage { pct: 50. },
            Percentage { pct: 50. },
            Percentage { pct: 25. },
            Overlap::Number { N: 10 },
        ];
        let nffts = [10, 10, 1024, 10];
        let expected_keep = [0, 5, 512, 2, 10];

        for ((expected, nfft), overlap) in expected_keep.iter().zip(nffts.iter()).zip(ol.iter()) {
            let settings = ApsSettingsBuilder::default()
                .nfft(*nfft)
                .fs(1.)
                .overlap(overlap.clone())
                .build()
                .expect("BUG: Settings cannot be build.");

            assert_eq!(settings.get_overlap_keep(), *expected);
        }
    }

    /// When the time constant is 1.0, every second the powers approximately
    /// halve. That is the subject of this test.
    #[test]
    fn test_expweighting() {
        let nfft = 48000;
        let fs = nfft as Flt;
        let tau = 2.;
        let settings = ApsSettingsBuilder::default()
            .fs(fs)
            .nfft(nfft)
            .overlap(Overlap::NoOverlap {})
            .mode(ApsMode::ExponentialWeighting { tau })
            .build()
            .expect("Settings cannot be empty");
        let overlap_keep = settings.get_overlap_keep();
        let mut aps = AvPowerSpectra::new(settings);
        assert_eq!(aps.overlap_keep, 0);

        let timedata_some = randNormal((nfft, 1));
        let timedata_zeros = Dmat::zeros((nfft, 1));

        // Clone here, as first_result reference is overwritten by subsequent
        // calls to compute_last.
        let first_result = aps.compute_last(timedata_some.view()).unwrap().clone();

        aps.compute_last(&timedata_zeros).unwrap();

        let last = aps.compute_last(&timedata_zeros).unwrap();

        let alpha = Flt::exp(-((nfft - overlap_keep) as Flt) / (fs * tau));

        for i in 0..nfft / 2 + 1 {
            assert_abs_diff_eq!(first_result.ap(0)[i] * alpha.powi(2), last.ap(0)[i]);
        }
        assert_eq!(aps.N, 3);
    }

    #[test]
    fn test_tf1() {
        let nfft = 4800;
        let mut timedata = randNormal((nfft, 1));
        timedata
            .push_column(timedata.column(0).mapv(|a| 2. * a).view())
            .unwrap();

        let settings = ApsSettingsBuilder::default()
            .fs(1.0)
            .nfft(nfft)
            .build()
            .unwrap();
        let mut aps = AvPowerSpectra::new(settings);
        if let Some(v) = aps.compute_last(&timedata) {
            let tf = v.tf(0, 1, None);
            assert_eq!((&tf - 2.0 * Cflt::ONE).sum().abs(), 0.0);
        } else {
            unreachable!()
        }
    }

    #[test]
    fn test_tf2() {
        let nfft = 4800;
        let mut timedata = randNormal((nfft, 1));
        timedata
            .push_column(timedata.column(0).mapv(|a| 2. * a).view())
            .unwrap();
        // Negative reference channel
        timedata
            .push_column(timedata.column(0).mapv(|a| -a).view())
            .unwrap();

        let settings = ApsSettingsBuilder::default()
            .fs(1.)
            .nfft(nfft)
            .build()
            .unwrap();
        let mut aps = AvPowerSpectra::new(settings);
        if let Some(v) = aps.compute_last(&timedata) {
            let tf = v.tf(0, 1, Some(2));
            assert_eq!((&tf - 2.0 * Cflt::ONE).sum().abs(), 0.0);
        } else {
            unreachable!()
        }
    }
    #[test]
    fn test_ap() {
        let nfft = 1024;
        let timedata = randNormal((150 * nfft, 1));
        let timedata_mean_square = (&timedata * &timedata).sum() / (timedata.len() as Flt);

        for wt in [
            Some(WindowType::Rect),
            Some(WindowType::Hann),
            Some(WindowType::Bartlett),
            Some(WindowType::Blackman),
            None,
        ] {
            let settings = ApsSettingsBuilder::default()
                .fs(1.0)
                .nfft(nfft)
                .windowType(wt.unwrap_or_default())
                .build()
                .unwrap();
            let mut aps = AvPowerSpectra::new(settings);
            if let Some(v) = aps.compute_last(&timedata) {
                let ap = v.ap(0);
                assert_abs_diff_eq!(ap.sum().abs(), timedata_mean_square, epsilon = 1e-2);
            } else {
                unreachable!()
            }
        }
    }
    #[test]
    #[should_panic]
    fn test_apssettings1() {
        let _ = ApsSettingsBuilder::default().build().unwrap();
    }

    #[test]
    fn test_apssettings2() {
        let _ = ApsSettingsBuilder::default()
            .nfft(2048)
            .fs(1.0)
            .build()
            .unwrap();
    }
}