lasprs 0.4.1

use super::timebuffer::TimeBuffer;
use super::CrossPowerSpecra;
use super::*;
use crate::config::*;
use anyhow::{bail, Result};

///  Provide the overlap of blocks for computing averaged (cross) power spectra.
///  Can be provided as a percentage of the block size, or as a number of
///  samples.
pub enum Overlap {
    /// Overlap specified as a percentage of the total FFT length. Value should
    /// be 0<=pct<100.
    Percentage(Flt),
    /// Number of samples to overlap
    Number(usize),
    /// No overlap at all, which is the same as Overlap::Number(0)
    NoOverlap,
}
impl Default for Overlap {
    fn default() -> Self {
        Overlap::Percentage(50.)
    }
}

/// Result from [AvPowerspectra.compute()]
pub enum ApsResult<'a> {
    /// Returns all intermediate results. Useful when evolution over time should
    /// be visible. I.e. in case of spectrograms, but also in case the
    /// converging process to the average should be made visible.
    AllIntermediateResults(Vec<CPS>),

    /// Give only last result back, the most recent value.
    OnlyLastResult(&'a CPS),

    /// No new data available. Nothing given back.
    None,
}

/// The 'mode' used in computing averaged power spectra. When providing data in
/// blocks to the [AvPowerSpectra] the resulting 'current estimate' responds
/// differently, depending on the model.
pub enum ApsMode {
    /// Averaged over all data provided. New averages can be created by calling
    /// `AvPowerSpectra::reset()`
    AllAveraging,
    /// In this mode, the `AvPowerSpectra` works a bit like a sound level meter,
    /// where new data is weighted with old data, and old data exponentially
    /// backs off. This mode only makes sense when `tau >> nfft/fs`
    ExponentialWeighting {
        /// Sampling frequency in [Hz], used for computing IIR filter coefficient.
        fs: Flt,
        /// Time weighting constant, follows convention of Sound Level Meters.
        /// Means the data is approximately low-pass filtered with a cut-off
        /// frequency f_c of s/tau ≅ 1 → f_c = (2 * pi * tau)^-1.
        tau: Flt,
    },
    /// Spectrogram mode. Only returns the latest estimate(s).
    Spectrogram,
}
impl Default for ApsMode {
    fn default() -> Self {
        ApsMode::AllAveraging
    }
}

/// 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.
///
pub struct AvPowerSpectra {
    // Power spectra estimator for single block
    ps: PowerSpectra,

    // The mode with which the AvPowerSpectra is computing.
    mode: ApsMode,

    // 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,

    // Current estimation of the power spectra
    cur_est: CPS,
}
impl AvPowerSpectra {
    /// The FFT Length of estimating (cross)power spectra
    pub fn nfft(&self) -> usize {
        self.ps.nfft()
    }
    /// Returns the amount of samples to `keep` in the time buffer when
    /// overlapping time segments using [TimeBuffer].
    fn get_overlap_keep(nfft: usize, overlap: Overlap) -> Result<usize> {
        let overlap_keep = match overlap {
            Overlap::Number(i) if i >= nfft => {
                bail!("Invalid overlap number of samples. Should be < nfft, which is {nfft}.")
            }
            // Keep 1 sample, if overlap is 1 sample etc.
            Overlap::Number(i) if i < nfft => i,

            // If overlap percentage is >= 100, or < 0.0 its an error
            Overlap::Percentage(p) if p >= 100. || p < 0.0 => {
                bail!("Invalid overlap percentage. Should be >= 0. And < 100.")
            }
            // If overlap percentage is 0, this gives
            Overlap::Percentage(p) => ((p * nfft as Flt) / 100.) as usize,
            Overlap::NoOverlap => 0,
            _ => unreachable!(),
        };
        if overlap_keep >= nfft {
            bail!("Computed overlap results in invalid number of overlap samples. Please make sure the FFT length is large enough, when high overlap percentages are required.");
        }
        Ok(overlap_keep)
    }

    /// 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 = CPS::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
    ///
    /// * `nfft` - FFT Length
    ///
    /// # Panics
    ///
    /// - When nfft is not even, or 0.
    /// 
    pub fn new_simple(nfft: usize) -> AvPowerSpectra {
        AvPowerSpectra::build(nfft, None, None, None).unwrap()
    }

    /// 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 build(
        nfft: usize,
        windowtype: Option<WindowType>,
        overlap: Option<Overlap>,
        mode: Option<ApsMode>,
    ) -> Result<AvPowerSpectra> {
        if nfft % 2 != 0 {
            bail!("NFFT should be even")
        }
        if nfft == 0 {
            bail!("Invalid NFFT")
        }
        let windowtype = windowtype.unwrap_or_default();
        let window = Window::new(windowtype, nfft);
        let mode = mode.unwrap_or_default();
        let overlap = overlap.unwrap_or_default();

        // Perform some checks on ApsMode
        if let ApsMode::ExponentialWeighting { fs, tau } = mode {
            if fs <= 0.0 {
                bail!("Invalid sampling frequency given as parameter");
            }
            if tau <= 0.0 {
                bail!("Invalid time weighting constant [s]. Should be > 0 if given.");
            }
        }

        let ps = PowerSpectra::newFromWindow(window);
        let overlap_keep = Self::get_overlap_keep(nfft, overlap)?;

        Ok(AvPowerSpectra {
            ps,
            overlap_keep,
            mode,
            N: 0,
            cur_est: CPS::default((0, 0, 0)),
            timebuf: TimeBuffer::new(),
        })
    }
    // Update result for single block
    fn update_singleblock<'a, T>(&mut self, timedata: T)
    where
        T: AsArray<'a, Flt, Ix2>,
    {
        let timeblock = timedata.into();
        let Cpsnew = self.ps.compute(&timeblock);
        // println!("Cpsnew: {:?}", Cpsnew[[0, 0, 0]]);

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

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

        // Apply operation based on mode
        match self.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 { fs, 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 / 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).
    /// * `giveInterMediateResults` - If true: returns a vector of all
    ///   intermediate results. Useful when plotting spectra over time.
    ///
    /// # Panics
    ///
    /// If timedata.ncols() does not match number of columns in already present
    /// data.
    pub fn compute<'a, 'b, T>(
        &'a mut self,
        timedata: T,
        giveInterMediateResults: bool,
    ) -> ApsResult<'a>
    where
        T: AsArray<'b, Flt, Ix2>,
    {
        // Push new data in the time buffer.
        self.timebuf.push(timedata);

        // Storage for the result
        let mut result = ApsResult::None;

        // 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);

            // We ha
            computed_single = true;
            if giveInterMediateResults {
                // Initialize with empty vector
                if let ApsResult::None = result {
                    result = ApsResult::AllIntermediateResults(Vec::new())
                }
            }

            // So
            if let ApsResult::AllIntermediateResults(v) = &mut result {
                v.push(self.cur_est.clone());
            }
        }

        if computed_single && !giveInterMediateResults {
            // We have computed it once, but we are not interested in
            // intermediate results. So we return a reference to the last data.
            return ApsResult::OnlyLastResult(&self.cur_est);
        }
        result
    }
    /// See [AvPowerSpectra](compute())
    pub fn compute_last<'a, T>(&mut self, timedata: T) -> ApsResult
    where
        T: AsArray<'a, Flt, Ix2>,
    {
        return self.compute(timedata, false);
    }
}

#[cfg(test)]
mod test {
    use approx::assert_abs_diff_eq;
    use ndarray_rand::rand_distr::Normal;
    use ndarray_rand::RandomExt;

    use super::CrossPowerSpecra;
    use crate::{config::*, ps::ApsResult};

    use super::{ApsMode, AvPowerSpectra, Overlap, WindowType, CPS};

    #[test]
    fn test_overlap_keep() {
        let nfft = 10;
        assert_eq!(
            AvPowerSpectra::get_overlap_keep(nfft, Overlap::Percentage(50.)).unwrap(),
            nfft / 2
        );
        let nfft = 11;
        assert_eq!(
            AvPowerSpectra::get_overlap_keep(nfft, Overlap::Percentage(50.)).unwrap(),
            nfft / 2
        );
        let nfft = 1024;
        assert_eq!(
            AvPowerSpectra::get_overlap_keep(nfft, Overlap::Percentage(25.)).unwrap(),
            nfft / 4
        );
        assert_eq!(
            AvPowerSpectra::get_overlap_keep(nfft, Overlap::Number(10)).unwrap(),
            10
        );
    }

    /// 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 mut aps = AvPowerSpectra::build(
            nfft,
            None,
            Some(Overlap::NoOverlap),
            Some(ApsMode::ExponentialWeighting { fs, tau }),
        )
        .unwrap();
        assert_eq!(aps.overlap_keep, 0);

        let distr = Normal::new(1.0, 1.0).unwrap();
        let timedata_some = Dmat::random((nfft, 1), distr);
        let timedata_zeros = Dmat::zeros((nfft, 1));

        let mut first_result = CPS::zeros((0, 0, 0));

        if let ApsResult::OnlyLastResult(v) = aps.compute_last(timedata_some.view()) {
            first_result = v.clone();
            // println!("{:?}", first_result.ap(0)[0]);
        } else {
            assert!(false, "Should return one value");
        }
        let overlap_keep = AvPowerSpectra::get_overlap_keep(nfft, Overlap::NoOverlap).unwrap();
        if let ApsResult::OnlyLastResult(_) = aps.compute_last(&timedata_zeros) {
            // Do nothing with it.
        } else {
            assert!(false, "Should return one value");
        }
        if let ApsResult::OnlyLastResult(v) = aps.compute_last(&timedata_zeros) {
            let alpha = Flt::exp(-((nfft - overlap_keep) as Flt) / (fs * tau));
            // let alpha: Flt = 1.;
            for i in 0..nfft / 2 + 1 {
                // println!("i={i}");
                assert_abs_diff_eq!(first_result.ap(0)[i] * alpha.powi(2), v.ap(0)[i]);
            }
        } else {
            assert!(false, "Should return one value");
        }
        assert_eq!(aps.N, 3);
    }

    #[test]
    fn test_tf1() {
        let nfft = 4800;
        let distr = Normal::new(1.0, 1.0).unwrap();
        let mut timedata = Dmat::random((nfft, 1), distr);
        timedata
            .push_column(timedata.column(0).mapv(|a| 2. * a).view())
            .unwrap();

        let mut aps = AvPowerSpectra::build(nfft, None, None, None).unwrap();
        if let ApsResult::OnlyLastResult(v) = aps.compute_last(&timedata) {
            let tf = v.tf(0, 1, None);
            assert_eq!((&tf - 2.0 * Cflt::ONE).sum().abs(), 0.0);
        } else {
            assert!(false);
        }
    }

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

        let mut aps = AvPowerSpectra::build(nfft, None, None, None).unwrap();
        if let ApsResult::OnlyLastResult(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 {
            assert!(false);
        }
    }
    #[test]
    fn test_ap() {
        let nfft = 1024;
        let distr = Normal::new(1.0, 1.0).unwrap();
        let timedata = Dmat::random((150 * nfft, 1), distr);
        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 mut aps = AvPowerSpectra::build(nfft, wt, None, None).unwrap();
            if let ApsResult::OnlyLastResult(v) = aps.compute_last(&timedata) {
                let ap = v.ap(0);
                assert_abs_diff_eq!((&ap).sum().abs(), timedata_mean_square, epsilon = 1e-2);
            } else {
                assert!(false);
            }
        }
    }
}