lasprs 0.6.7

use super::siggenchannel::SiggenChannelConfig;
use super::source::{self, *};
use super::sweep::SweepType;
use super::SiggenCommand;
use crate::config::*;
use crate::filter::Filter;
use anyhow::{bail, Result};
use dasp_sample::{FromSample, Sample};
use rayon::prelude::*;
use std::fmt::Debug;
use std::iter::ExactSizeIterator;
use std::slice::IterMut;

/// Dummy sampling frequency to be filled in when the sampling frequency is
/// still unknown at the point in time.
pub const DUMMY_SAMPLING_FREQ: Flt = 48000.;

/// Multiple channel signal generator. Able to create (acoustic) output signals. See above example on how to use.
/// Typical signal that can be created are:
///
/// * [Siggen::newWhiteNoise]
/// * [Siggen::newSine]
/// * [Siggen::newSilence]
///
#[derive(Clone, Debug)]
#[cfg_attr(feature = "python-bindings", pyclass)]
pub struct Siggen {
    // Sampling frequency in Hz
    fs: Option<Flt>,

    // The source dynamic signal. Noise, a sine wave, sweep, etc
    source: Source,

    // Channel configuration for each output channel
    channels: Vec<SiggenChannelConfig>,

    // Temporary source signal buffer
    source_buf: Vec<Flt>,

    // Output buffers (for filtered source signal)
    chout_buf: Vec<Vec<Flt>>,
}
#[cfg(feature = "python-bindings")]
#[cfg_attr(feature = "python-bindings", pymethods)]
impl Siggen {
    #[new]
    fn new_py() -> Self {
        Siggen::new(1, Source::newSilence())
    }
}

impl Siggen {
    /// Create a new signal generator with an arbitrary source.
    /// # Args
    ///
    /// - `nchannels` - The number of channels to output
    /// - `source` - Source that generates the signal
    pub fn new(nchannels: usize, source: Source) -> Siggen {
        Siggen {
            fs: None,
            source,
            channels: vec![SiggenChannelConfig::new(); nchannels],
            source_buf: vec![],
            chout_buf: vec![],
        }
    }
    /// Create sine sweep signal generator
    ///
    /// # Args
    ///
    /// - `fs` - Sample rate \[Hz\]
    /// - `nchannels`: The number of channels to output
    /// - `fl` - Lower frequency \[Hz\]
    /// - `fu` - Upper frequency \[Hz\]
    /// - `sweep_time` - The duration of a single sweep \[s\]
    /// - `quiet_time` - Time of silence after one sweep and start of the next \[s\]
    /// - `sweep_type` - The type of the sweep, see [SweepType].
    pub fn newSweep(
        fs: Flt,
        nchannels: usize,
        fl: Flt,
        fu: Flt,
        sweep_time: Flt,
        quiet_time: Flt,
        sweep_type: SweepType,
    ) -> Result<Self> {
        let source = Source::newSweep(fs, fl, fu, sweep_time, quiet_time, sweep_type)?;
        Ok(Self::new(nchannels, source))
    }
    /// Create a sine wave signal generator
    ///
    /// # Args
    ///
    /// - `fs` - Sampling frequency \[Hz\]
    /// - `nchannels`: The number of channels to output
    /// * `freq` - Frequency of the sine wave in \[Hz\]
    pub fn newSine(fs: Flt, nchannels: usize, freq: Flt) -> Result<Siggen> {
        Ok(Siggen::new(nchannels, Source::newSine(fs, freq)?))
    }

    /// Silence: create a signal generator that does not output any dynamic
    /// signal at all.
    /// # Args
    ///
    /// - `fs` - Sampling frequency \[Hz\]
    /// - `nchannels` - The number of channels to output
    pub fn newSilence(_fs: Flt, nchannels: usize) -> Siggen {
        Siggen::new(nchannels, Source::newSilence())
    }

    /// Create a white noise signal generator.
    ///
    /// # Args
    ///
    /// - `fs` - Sampling frequency \[Hz\]
    /// - `nchannels` - The number of channels to output
    pub fn newWhiteNoise(fs: Flt, nchannels: usize, interrupt_period: Option<Flt>) -> Siggen {
        Siggen::new(nchannels, Source::newWhiteNoise(fs, interrupt_period))
    }

    /// Returns the number of channels this signal generator is generating for.
    pub fn nchannels(&self) -> usize {
        self.channels.len()
    }

    /// Apply command to current signal generator to change its state.
    pub fn applyCommand(&mut self, msg: SiggenCommand) -> Result<()> {
        match msg {
            SiggenCommand::ChangeSource { src } => {
                self.source = src;

                self.reset(self.fs.unwrap_or(48e3));
                Ok(())
            }
            SiggenCommand::ResetSiggen { fs } => {
                self.reset(fs);
                Ok(())
            }
            SiggenCommand::SetMuteAllChannels { mute } => {
                self.setAllMute(mute);
                Ok(())
            }
            SiggenCommand::SetMuteChannel { ch, mute } => {
                if ch > self.channels.len() {
                    bail!("Invalid channel index: {ch}");
                }
                self.channels[ch].setMute(mute);
                Ok(())
            }
            SiggenCommand::SetAllGains { g } => {
                self.setAllGains(g);
                Ok(())
            }
        }
    }
    /// Set gains of all channels in signal generator to the same value
    ///
    /// # Args
    ///
    /// * g: New gain value
    pub fn setAllGains(&mut self, g: Flt) {
        self.channels.iter_mut().for_each(|set| set.setGain(g))
    }

    /// Set the number of channels to generate a signal for. Truncates the
    /// output in case the value before calling this method is too little.
    /// Appends new channel configs in case to little is available.
    ///
    /// * nch: The new required number of channels
    pub fn setNChannels(&mut self, nch: usize) {
        self.channels.truncate(nch);

        while self.channels.len() < nch {
            self.channels.push(SiggenChannelConfig::new());
        }
    }

    /// Set the DC offset for all channels
    pub fn setDCOffset(&mut self, dc: &[Flt]) {
        self.channels.iter_mut().zip(dc).for_each(|(ch, dc)| {
            ch.DCOffset = *dc;
        });
    }

    /// Creates *interleaved* output signal
    pub fn genSignal<T>(&mut self, out: &mut [T])
    where
        T: Sample + FromSample<Flt> + Debug,
        Flt: Sample,
    {
        let nch = self.nchannels();
        let nsamples: usize = out.len() / nch;
        assert!(out.len() % self.nchannels() == 0);

        // Create source signal
        self.source_buf.resize(nsamples, 0.0);
        self.source
            .genSignal_unscaled(&mut self.source_buf.iter_mut());
        // println!("Source signal: {:?}", self.source_buf);

        // Write output while casted to the correct type
        // Iterate over each channel, and counter
        self.chout_buf.resize(nch, vec![]);

        for (channelno, (channel, chout)) in self
            .channels
            .iter_mut()
            .zip(self.chout_buf.iter_mut())
            .enumerate()
        {
            chout.resize(nsamples, 0.0);

            // Create output signal, overwrite chout
            channel.genSignal(&self.source_buf, chout);
            // println!("Channel: {}, {:?}", channelno, chout);

            let out_iterator = out.iter_mut().skip(channelno).step_by(nch);
            out_iterator.zip(chout).for_each(|(out, chin)| {
                *out = chin.to_sample();
            });
        }
        // println!("{:?}", out);
    }

    /// Reset signal generator. Applies any kind of cleanup necessary.
    ///
    /// Args
    ///
    /// * fs: (New) Sampling frequency \[Hz\]
    ///
    pub fn reset(&mut self, fs: Flt) {
        self.fs = Some(fs);
        self.source.reset(fs);
        self.channels.iter_mut().for_each(|x| x.reset(fs))
    }
    /// Mute / unmute all channels at once
    pub fn setAllMute(&mut self, mute: bool) {
        self.channels.iter_mut().for_each(|s| {
            s.setMute(mute);
        });
    }

    /// Mute / unmute individual channels. Array of bools should have same size
    /// as number of channels in signal generator.
    pub fn setMute(&mut self, mute: &[bool]) {
        assert!(mute.len() == self.nchannels());
        self.channels.iter_mut().zip(mute).for_each(|(s, m)| {
            s.setMute(*m);
        });
    }
}

#[cfg(test)]
mod test {
    use approx::assert_abs_diff_eq;

    use super::*;
    use crate::Flt;

    #[test]
    fn test_whitenoise() {
        // This code is just to check syntax. We should really be listening to these outputs.
        let mut t = [0.0; 10];
        Siggen::newWhiteNoise(1., 1, None).genSignal(&mut t);
        // println!("{:?}", &t);
    }

    #[test]
    fn test_sine() {
        // This code is just to check syntax. We should really be listening to
        // these outputs.
        const N: usize = 10000;
        let mut s1 = [0.0; N];
        let mut s2 = [0.0; N];
        let mut siggen = Siggen::newSine(10., 1, 1.0).unwrap();

        siggen.reset(10.0);
        siggen.setAllMute(false);
        siggen.genSignal(&mut s1);
        siggen.reset(10.0);
        siggen.genSignal(&mut s2);

        let absdiff = s1
            .iter()
            .zip(s2.iter())
            .map(|(s1, s2)| Flt::abs(*s1 - *s2))
            .sum::<Flt>();
        assert_abs_diff_eq!(absdiff, 0., epsilon = Flt::EPSILON * 100.);
    }

    #[test]
    fn test_sine2() {
        // Test if channels are properly separated etc. Check if RMS is correct
        // for amplitude = 1.0.
        const fs: Flt = 10.0;
        // Number of samples per channel
        const Nframes: usize = 10000;
        const Nch: usize = 2;
        let mut signal = [0.0; Nch * Nframes];
        let mut siggen = Siggen::newSine(fs, Nch, 1.0).unwrap();

        siggen.reset(fs);
        siggen.setMute(&[false, true]);
        // siggen.channels[0].DCOffset = 0.1;

        // Split off in two terms, see if this works properly
        siggen.genSignal(&mut signal[..Nframes / 2]);
        siggen.genSignal(&mut signal[Nframes / 2..]);

        // Mean square of the signal
        let ms1 = signal.iter().step_by(2).map(|s1| *s1 * *s1).sum::<Flt>() / (Nframes as Flt);
        println!("ms1: {}", ms1);

        let ms2 = signal
            .iter()
            .skip(1)
            .step_by(2)
            .map(|s1| *s1 * *s1)
            .sum::<Flt>()
            / (Nframes as Flt);

        assert_abs_diff_eq!(Flt::abs(ms1 - 0.5), 0., epsilon = Flt::EPSILON * 1e3);
        assert_eq!(ms2, 0.0);
    }

    // A small test to learn a bit about sample types and conversion. This
    // is the thing we want.
    #[test]
    fn test_sample() {
        assert_eq!((0.5f32).to_sample::<i8>(), 64);
        assert_eq!((1.0f32).to_sample::<i8>(), 127);
        assert_eq!(-(1.0f32).to_sample::<i8>(), -127);
        assert_eq!((1.0f32).to_sample::<i16>(), i16::MAX);
    }
}