lasprs 0.8.0

use crate::config::*;
use derive_builder::Builder;
use itertools::Itertools;
use ndarray::ArrayView1;
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
use rayon::prelude::*;
use smallvec::SmallVec;

use super::{slmsettings::SLMSettings, SLM_MAX_CHANNELS};
use crate::{config::*, filter::Filter};
use crate::{Biquad, Dcol, Flt, FreqWeighting, PoleOrZero, SeriesBiquad, ZPKModel};
#[derive(Debug, Clone)]
struct SLMChannel {
    // Number of samples processed after last run() is called.
    N: usize,

    // Square of the reference level, i.e. (20 microPascal)^2
    Lrefsq: Flt,

    // Statistics to store
    stat: SLMStat,

    // The bandpass filter for this channel, if any
    bp: SeriesBiquad,

    // Time weighting filter
    twfilter: TwFilter,
}
impl SLMChannel {
    fn run(&mut self, prefiltered: &[Flt], provide_output: bool) -> Option<Vec<Flt>> {
        // Do not continue below if there is no data to run on
        if prefiltered.is_empty() {
            if provide_output {
                return Some(vec![]);
            } else {
                return None;
            }
        }

        assert!(!prefiltered.is_empty());
        // Simple level computer for power inputs
        let level_fun = |a| 10. * Flt::log10(a / self.Lrefsq);

        let mut tmp = self.bp.filter(prefiltered);
        let mut N = self.N;

        // Filtered squared
        let mut filtered_squared = {
            let mut tmp_view = ArrayViewMut1::from(&mut tmp);
            tmp_view.mapv_inplace(|a| a * a);
            tmp_view
        };

        // Compute and update Lpk and Leq
        filtered_squared.for_each(|sample_pwr| {
            let new_pk = sample_pwr.abs();
            if new_pk > self.stat.Ppk {
                self.stat.Ppk = new_pk;
            }
            // Update equivalent level
            self.stat.Peq = (self.stat.Peq * N as Flt + sample_pwr) / (N as Flt + 1.);
            N += 1;
        });

        // Apply time weighting to filtered squared signal
        filtered_squared.mapv_inplace(|s| self.twfilter.filter(s));

        // Update max signal power gotten so far
        let time_weighted = &mut filtered_squared;
        time_weighted.for_each(|val| {
            if *val > self.stat.Pmax {
                self.stat.Pmax = *val;
            }
        });

        // Update last signal power coming from SLM
        self.stat.Pt_last = *filtered_squared.last().unwrap();
        // Convert output to levels
        filtered_squared.mapv_inplace(level_fun);

        self.N += prefiltered.len();
        if provide_output {
            Some(tmp)
        } else {
            None
        }
    }
}

/// Sound Level Meter
#[cfg_attr(feature = "python-bindings", pyclass)]
#[derive(Debug, Clone)]
pub struct SLM {
    prefilter: SeriesBiquad,
    channels: SmallVec<[SLMChannel; SLM_MAX_CHANNELS]>,
}

impl SLM {
    /// Create new Sound Level Meter from given settings
    pub fn new(settings: SLMSettings) -> Self {
        let fs = settings.fs;

        // Generate rectifier filter(s)
        let poles = settings.timeWeighting.getLowpassPoles();
        let alpha_up = Flt::exp(poles.0 / fs);
        let alpha_down = poles.1.map(|p| Flt::exp(p / fs));
        // dbg!(alpha_up);
        // dbg!(alpha_down);

        let twfilter = TwFilter::new(alpha_up, alpha_down);
        let prefilter = ZPKModel::freqWeightingFilter(settings.freqWeighting).bilinear(fs);
        let channels = settings
            .filterDescriptors
            .iter()
            .map(|descriptor| {
                // Generate bandpass filter
                let bp = descriptor.genFilter().bilinear(fs);
                // Initalize statistics with defaults
                let stat = SLMStat::default();

                SLMChannel {
                    stat,
                    bp,
                    twfilter,
                    Lrefsq: settings.Lref.powi(2),
                    N: 0,
                }
            })
            .collect();
        SLM {
            prefilter,
            channels,
        }
    }
    /// Push new time data through sound level meter. Returns L(t) data for each
    /// channel. Updates the computed statistics and optionally outputs levels
    /// vs time if flag `provide_output` is set to `true`. Note that at the end
    /// of the block, the `L(t)` can also be obtained by calling [SLM::Ltlast].
    ///
    /// # Args
    ///
    /// - `td`: Time data
    /// - `provide_output` - Set this to true to give intermediate output data
    pub fn run(&mut self, td: &[Flt], provide_output: bool) -> Option<Vec<Vec<Flt>>> {
        if td.is_empty() {
            // No input, so no output
            return None;
        }
        // Pre-filter, e.g. apply A-weighting to the data
        let prefiltered = self.prefilter.filter(td);

        let ch_iter = self.channels.par_iter_mut().map(|ch| {
            // Loop over each bandpass channel, and compute SLM for each channel
            // in parallel. The iterator itself here does not do anything, only
            // the collect / for_each statements below.
            ch.run(&prefiltered, provide_output)
        });
        if provide_output {
            // Collect and unwrap
            let Lt: Vec<_> = ch_iter
                .map(|a| a.expect("No output provided for ch"))
                .collect();
            Some(Lt)
        } else {
            // Just consume the iterator, but do not return anythin
            ch_iter.for_each(|_| {});
            None
        }
    }

    /// Number of channels in SLM
    pub fn nch(&self) -> usize {
        self.channels.len()
    }

    fn levels_from<T>(&self, stat_returner: T) -> Dcol
    where
        T: Fn(&SLMChannel) -> Flt,
    {
        Dcol::from_iter(
            self.channels
                .iter()
                .map(|ch| 10. * Flt::log10(stat_returner(ch) / ch.Lrefsq)),
        )
    }

    /// Get max levels for each channel
    pub fn Lmax(&self) -> Dcol {
        self.levels_from(|ch| ch.stat.Pmax)
    }
    /// Get peak levels for each channel
    pub fn Lpk(&self) -> Dcol {
        self.levels_from(|ch| ch.stat.Ppk)
    }

    /// Get equivalent levels for each channel
    pub fn Leq(&self) -> Dcol {
        self.levels_from(|ch| ch.stat.Peq)
    }
    /// Get last value of level vs time
    pub fn Ltlast(&self) -> Dcol {
        self.levels_from(|ch| ch.stat.Pt_last)
    }
}

#[cfg(feature = "python-bindings")]
#[cfg_attr(feature = "python-bindings", pymethods)]
impl SLM {
    #[new]
    fn new_py(settings: SLMSettings) -> SLM {
        SLM::new(settings)
    }

    #[pyo3(name = "run", signature=(dat, provide_output=true))]
    fn run_py<'py>(
        &mut self,
        py: Python<'py>,
        dat: PyArrayLike1<Flt>,
        provide_output: bool,
    ) -> Option<Vec<Bound<'py, PyArray1<Flt>>>> {
        if let Some(res) = self.run(dat.as_array().as_slice()?, provide_output) {
            let vec_py_iter = res.into_iter().map(|v| PyArray1::from_vec(py, v));
            let vec_py = vec_py_iter.collect::<Vec<Bound<'py, PyArray1<Flt>>>>();
            Some(vec_py)
        } else {
            None
        }
    }

    #[pyo3(name = "Lmax")]
    fn Lmax_py<'py>(&self, py: Python<'py>) -> PyArr1Flt<'py> {
        PyArray1::from_array(py, &self.Lmax())
    }
    #[pyo3(name = "Leq")]
    fn Leq_py<'py>(&self, py: Python<'py>) -> PyArr1Flt<'py> {
        PyArray1::from_array(py, &self.Leq())
    }
    #[pyo3(name = "Lpk")]
    fn Lpk_py<'py>(&self, py: Python<'py>) -> PyArr1Flt<'py> {
        PyArray1::from_array(py, &self.Lpk())
    }
    #[pyo3(name = "Ltlast")]
    fn Ltlast_py<'py>(&self, py: Python<'py>) -> PyArr1Flt<'py> {
        PyArray1::from_array(py, &self.Ltlast())
    }
}

#[derive(Debug, Clone, Default)]
/// Quantities defined as powers, i.e. square of amplitude
struct SLMStat {
    // Max signal power
    Pmax: Flt,
    // Peak signal power
    Ppk: Flt,
    // Equivalent signal power
    Peq: Flt,

    // Last obtained signal power, after last time run() is called.
    Pt_last: Flt,
}

#[derive(Clone, Copy, Debug)]
struct F {
    alpha: Flt,
    state: Flt,
}
// Time weighting filter
#[derive(Clone, Copy, Debug)]
struct TwFilter {
    // Filter constant (see lasp_doc) for time weighting and state of previous sample.
    filter_state_up: F,

    // Filter constant (see lasp_doc) for time weighting (only there for
    // asymmetric filters (i.e. impulse time weighting)
    filter_state_down: Option<F>,
}
impl TwFilter {
    fn new(alpha_up: Flt, alpha_down: Option<Flt>) -> TwFilter {
        let filter_state_down = alpha_down.map(|alpha_down| F {
                alpha: alpha_down,
                state: 0.,
            });
        TwFilter {
            filter_state_up: F {
                alpha: alpha_up,
                state: 0.,
            },
            filter_state_down,
        }
    }

    #[inline]
    fn filter(&mut self, input: Flt) -> Flt {
        if self.filter_state_down.is_some() {
            self.filter_asymmetric(input)
        } else {
            self.filter_symmetric(input)
        }
    }
    #[inline]
    fn filter_symmetric(&mut self, input: Flt) -> Flt {
        assert!(
            self.filter_state_down.is_none(),
            "Call this only for symmetric filters"
        );
        let F { alpha, state } = &mut self.filter_state_up;
        let alpha = *alpha;
        let output = (1. - alpha) * input + alpha * *state;
        *state = output;
        output
    }
    #[inline]
    fn filter_asymmetric(&mut self, input: Flt) -> Flt {
        let F {
            alpha: alpha_up,
            state: state_up,
        } = &mut self.filter_state_up;
        let alpha_up = *alpha_up;

        let filter_state_down = self
            .filter_state_down
            .as_mut()
            .expect("Call this only for asymmetric filters");
        let F {
            alpha: alpha_down,
            state: state_down,
        } = filter_state_down;
        let alpha_down = *alpha_down;

        let pu = (1. - alpha_up) * input + alpha_up * *state_up;
        *state_up = pu;

        let pIn = alpha_down * *state_down;

        
        if pu > pIn {
            *state_down = pu;
            pu
        } else {
            *state_down = pIn;
            pIn
        }
    }
}

#[cfg(test)]
mod test {
    use crate::{
        siggen::Siggen,
        slm::{SLMSettingsBuilder, TimeWeighting},
        Flt, FreqWeighting, StandardFilterDescriptor,
    };

    use super::SLM;

    #[test]
    fn test_slm1() {
        const fs: Flt = 48e3;
        const N: usize = (fs / 8.) as usize;

        let desc = StandardFilterDescriptor::Overall().unwrap();

        let settings = SLMSettingsBuilder::default()
            .fs(fs)
            .timeWeighting(TimeWeighting::Fast {})
            .freqWeighting(FreqWeighting::Z)
            .filterDescriptors(&[desc])
            .build()
            .unwrap();

        let mut siggen = Siggen::newSine(1., 1, 1000.).unwrap();
        siggen.setAllMute(false);
        siggen.reset(fs);
        let mut data = vec![0.; N];
        siggen.genSignal(&mut data);

        let mut slm = SLM::new(settings);
        // println!("{slm:#?}");
        let res = slm.run(&data, true).unwrap();
        let res = &res[0];
        println!("{slm:#?}");
        println!("{:#?}", &res[res.len() - 100..]);
    }
}