lasprs 0.14.1

use super::source::{Source, SourceImpl};
use super::*;
use crate::filter::{Equalizer, FilterGenerator, FilterMethods};
use crate::siggen::siggenchannel::SiggenChannelConfig;
use crate::*;
use crossbeam::channel::Sender;
use dasp_sample::{FromSample, Sample};
use ndarray::ArcArray2;
use rayon::prelude::*;
use snafu::prelude::*;
use std::fmt::Debug;
use std::iter::ExactSizeIterator;
use std::slice::IterMut;
use std::sync::Arc;

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

/// Multiple channel signal generator. Able to create (acoustic) output test
/// signals. See above example on how to use.
///
#[derive(Debug, Clone)]
pub(crate) struct Siggen {
    /// The source dynamic signal. Noise, a sine wave, sweep, etc
    srcdesc: SourceDescriptor,

    /// The actual source, instantiated from the source descriptor and sampling
    /// frequency
    source: Option<Source>,

    /// Channel configuration for each output channel, and the applied
    /// equalizer, if any.
    channels: Vec<(SiggenChannelConfig, Option<Equalizer>)>,

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

    /// Output buffers (for filtered source signal)
    chout_buf: Vec<Vec<Flt>>,
}

impl Siggen {
    /// Create a new signal generator with an arbitrary source.
    /// # Args
    ///
    /// - `nchannels` - The number of channels to output
    /// - `srcdesc` - Sourcedescriptor for the source that generates the signal
    pub fn new(nchannels: usize, srcdesc: SourceDescriptor) -> Siggen {
        Siggen {
            srcdesc,
            source: None,
            channels: vec![(SiggenChannelConfig::new(), None); nchannels],
            source_buf: vec![],
            chout_buf: vec![],
        }
    }

    /// 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 } => {
                if let Some(ref source) = self.source {
                    let fs = source.fs();
                    let newsrc = Source::new(&src, fs)?;
                    self.srcdesc = src;
                    self.source = Some(newsrc);
                } else {
                    self.srcdesc = src;
                }
                Ok(())
            }
            SiggenCommand::SetEqualizer { channels, eq } => {
                ensure!(
                    !channels.is_empty(),
                    InvalidParameterSnafu {
                        param: "channels",
                        criterion: "must not be empty"
                    }
                );
                // Unwrap is allowed here, above is already checked that
                // channels is not empty
                let max_ch = *channels.iter().max().unwrap();
                ensure!(
                    max_ch < self.channels.len(),
                    ChannelOutOfRangeSnafu {
                        channel: max_ch,
                        nchannels: self.channels.len()
                    }
                );
                if let Some(src) = &self.source {
                    // When we get here, a source is present, configured with a
                    // sample rate
                    let fs = src.fs();
                    if let Some(eq) = eq {
                        // Install filter based on EQ
                        let filter = eq.genFilter(fs).context(EqualizerSnafu {})?;
                        for chidx in channels {
                            // Yeah, we could use itertools::repeat_n here to
                            // avoid the clone for the last item, but we are
                            // lazy and cloning an eq and filter is probably not
                            // that expensive at all.
                            self.channels[chidx].0.setPreFilter(Some(filter.clone()));
                            self.channels[chidx].1 = Some(eq.clone());
                        }
                    } else {
                        // Disable existing EQ for given channels
                        for chidx in channels {
                            self.channels[chidx].0.setPreFilter(None);
                            self.channels[chidx].1 = None
                        }
                    }
                } else {
                    // No checks, stream is not running so signal generator does
                    // not (yet) use this Eq.
                    for chidx in channels {
                        self.channels[chidx].1 = eq.clone();
                    }
                }

                Ok(())
            }
            SiggenCommand::ResetSiggen { fs } => {
                self.reset(fs)?;
                Ok(())
            }
            SiggenCommand::SetMuteAllChannels { mute } => {
                self.setAllMute(mute);
                Ok(())
            }
            SiggenCommand::SetMuteChannel { ch, mute } => {
                ensure!(
                    ch < self.channels.len(),
                    ChannelOutOfRangeSnafu {
                        channel: ch,
                        nchannels: self.channels.len()
                    }
                );
                self.channels[ch].0.setMute(mute);
                Ok(())
            }
            SiggenCommand::SetAllGains { g } => {
                self.setAllGains(g);
                Ok(())
            }
            SiggenCommand::SetGain { ch, g } => {
                ensure!(
                    ch < self.channels.len(),
                    ChannelOutOfRangeSnafu {
                        channel: ch,
                        nchannels: self.channels.len()
                    }
                );
                self.channels[ch].0.setGain(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: Positive) {
        self.channels.iter_mut().for_each(|set| set.0.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(), None));
        }
    }

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

    /// Creates *interleaved* output signal, which is written to the provided buffer.
    pub fn genSignal<T>(&mut self, out: &mut [T])
    where
        T: Sample + FromSample<Flt> + Debug,
        Flt: Sample,
    {
        // Number of channels to generate signal for
        let nch = self.nchannels();
        // Number of samples to generate for each channel
        let nsamples: usize = out.len() / nch;
        assert!(out.len().is_multiple_of(self.nchannels()));

        // Create source signal
        self.source_buf.resize(nsamples, 0.0);
        if let Some(ref mut source) = self.source {
            source.genSignal_unscaled(&mut self.source_buf.iter_mut());
        } else {
            // No source signal.
            self.source_buf.fill(0.0);
        }

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

        // A bit of overhead here that the data is initialized as zeros, where
        // it is later overwritten. It is however really hard to code this such
        // that  this array is filled from an iterator.
        let mut out_monitor = ArcArray2::zeros((nsamples, nch).f());

        for (channelno, ((channel, chout_buf), mut out_mon_ch)) in self
            .channels
            .iter_mut()
            .zip(self.chout_buf.iter_mut())
            .zip(out_monitor.columns_mut())
            .enumerate()
        {
            // Make sure the size of the buffer is matching. This method does
            // nothing when the size is already correct.
            chout_buf.resize(nsamples, 0.0);

            // Create output signal, overwrite chout buffer. Here it is still
            // floating point
            channel.0.genSignal(&self.source_buf, chout_buf);

            // Copy over the channel into the monitor buffer
            out_mon_ch
                .as_slice_mut()
                .unwrap()
                .copy_from_slice(chout_buf);

            // Output iterator where the output needs to be copied (finally), in
            // `Sample` form.
            let out_iterator = out.iter_mut().skip(channelno).step_by(nch);

            // Copy over and convert to sample
            out_iterator.zip(chout_buf).for_each(|(out, chin)| {
                *out = chin.to_sample();
            });
        }
    }

    /// Reset signal generator. Applies any kind of cleanup necessary.
    ///
    /// Args
    ///
    /// * fs: (New) Sampling frequency \[Hz\]
    ///
    pub fn reset(&mut self, fs: StrictlyPositive) -> Result<()> {
        let source = Source::new(&self.srcdesc, fs)?;
        self.source = Some(source);
        for ch in self.channels.iter_mut() {
            if let Some(ref eq) = ch.1 {
                let filter = eq.genFilter(fs).context(EqualizerSnafu {})?;
                ch.0.setPreFilter(Some(filter));
                ch.0.reset(fs);
            }
        }
        Ok(())
    }
    /// Mute / unmute all channels at once
    pub fn setAllMute(&mut self, mute: bool) {
        self.channels.iter_mut().for_each(|s| {
            s.0.setMute(mute);
        });
    }

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

    /// Create a new signal generator with white noise source.
    ///
    /// # Arguments
    ///
    /// * `nchannels`: Number of channels in signal generator.
    /// * `fs`: Sampling frequency \[Hz\].
    ///
    /// # Returns
    ///
    /// A new signal generator with white noise source.
    ///
    /// # Errors
    ///
    /// Returns an error if the sampling frequency is not positive.
    #[allow(dead_code)]
    pub fn newWhiteNoise(nchannels: usize, fs: Flt) -> Result<Self> {
        let mut siggen = Siggen::new(
            nchannels,
            SourceDescriptor::WhiteNoise {
                interrupt_time: None,
            },
        );
        siggen.reset(fs.try_into().context(ParameterOutOfRangeSnafu {
            parameter: "Sampling frequency",
        })?)?;
        Ok(siggen)
    }

    /// Create a new signal generator with sine wave source.
    ///
    /// # Arguments
    ///
    /// * `nchannels`: Number of channels in signal generator.
    /// * `fs`: Sampling frequency \[Hz\].
    /// * `freq`: Frequency of sine wave \[Hz\].
    ///
    /// # Returns
    ///
    /// A new signal generator with sine wave source.
    ///
    /// # Errors
    ///
    /// Returns an error if the sampling frequency or frequency is not positive,
    /// or if the sine frequency does not work with the sampling frequency.
    #[allow(dead_code)]
    pub fn newSine(nchannels: usize, fs: Flt, freq: Flt) -> Result<Self> {
        let mut siggen = Siggen::new(
            nchannels,
            SourceDescriptor::Sine {
                frequency: freq.try_into().context(ParameterOutOfRangeSnafu {
                    parameter: "Frequency",
                })?,
            },
        );
        siggen.reset(fs.try_into().context(ParameterOutOfRangeSnafu {
            parameter: "Sampling frequency",
        })?)?;
        Ok(siggen)
    }
}

#[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];
        let mut siggen = Siggen::newWhiteNoise(1, 1.).unwrap();
        siggen.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 fs = 10.;
        let mut siggen = Siggen::newSine(1, fs, 1.).unwrap();

        siggen.setAllMute(false);
        siggen.genSignal(&mut s1);
        siggen.reset(fs.try_into().unwrap()).unwrap();
        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.

        // 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(Nch, 10.0, 1.).unwrap();

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