lasprs 0.14.1

//! Sweep signal generation code
use std::ops::Deref;

use super::{source::SourceImpl, *};
use crate::*;
use serde::{Deserialize, Serialize};
use snafu::prelude::*;
use strum::{EnumMessage, IntoEnumIterator};
use strum_macros::{Display, EnumIter, EnumMessage};

use crate::*;
/// Number of iterations in converging Newton-Raphson method, required in the
/// function to make a continuous sweep, that runs both forward and backward in
/// such a way as to be C1-coninuous.
const NITER_NEWTON: usize = 20;

/// Maximum length of a sweep signal in seconds.
const SWEEP_MAX_LENGTH_S: i32 = 61;

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

/// Enumerator representing the type of sweep source to create. Used as
/// parameter in [SourceDescriptor::Sweep].
#[cfg_attr(
    feature = "python-bindings",
    gen_stub_pyclass_enum,
    pyclass(eq, eq_int, from_py_object)
)]
#[derive(Debug, PartialEq, Clone, Display, EnumMessage, Serialize, Deserialize, EnumIter)]
pub enum SweepType {
    /// Forward only logarithmic sweep, repeats itself
    #[strum(message = "Forward logarithmic")]
    ForwardLog,
    /// Reverse only logarithmic sweep, repeats itself
    #[strum(message = "Backward logarithmic")]
    BackwardLog,
    /// Continuous logarithmic sweep, repeats itself
    #[strum(message = "Continuous logarithmic")]
    ContinuousLog,

    /// Forward only linear sweep, repeats itself
    #[strum(message = "Forward linear")]
    ForwardLin,
    /// Reverse only linear sweep, repeats itself
    #[strum(message = "Backward linear")]
    BackwardLin,
    /// Continuous linear sweep, repeats itself
    #[strum(message = "Continuous linear")]
    ContinuousLin,
}

#[cfg(feature = "python-bindings")]
#[cfg_attr(feature = "python-bindings", gen_stub_pymethods, pymethods)]
impl SweepType {
    #[staticmethod]
    fn all() -> Vec<SweepType> {
        Self::iter().collect()
    }
    fn __str__(&self) -> String {
        self.get_message().unwrap().into()
    }
}

/// Settings for the sweep signal generator, used to configure the signal
/// generator for outputing a sweep signal.
#[cfg_attr(feature = "python-bindings", gen_stub_pyclass, pyclass(from_py_object))]
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SweepSettings {
    /// Sweep type (forward, backward, continuous)
    sweep_type: SweepType,
    /// Start frequency in Hz
    fl: Bounded<0, 100_0001>,
    /// End frequency in Hz
    fu: StrictBounded<0, 100_0000>,
    /// Sweep time in seconds, larger than 0 and smaller than SWEEP_MAX_LENGTH_S
    /// s
    sweep_time: StrictBounded<0, SWEEP_MAX_LENGTH_S>,
    /// Optional quiet time in seconds
    quiet_time: Option<StrictBounded<0, 101>>,

    /// Optional amplitude modulation.
    /// - First index: frequency in Hz
    /// - Second index: linear amplitude value.
    amplitude_modulation: Option<Vec<(Positive, Positive)>>,
}
impl SweepSettings {
    /// Create a new sweep settings instance.
    ///
    /// # Arguments
    /// * `fl` - Lower frequency in Hz
    /// * `fu` - Upper frequency in Hz
    /// * `sweep_time` - Sweep time in seconds
    /// * `sweep_type` - Sweep type (forward, backward, continuous)
    /// * `quiet_time` - Optional quiet time in seconds
    /// * `amplitude_modulation` - Optional amplitude modulation
    ///
    /// # Returns
    /// A new `SweepSettings` instance
    pub fn new(
        fl: Flt,
        fu: Flt,
        sweep_time: Flt,
        sweep_type: SweepType,
        quiet_time: Option<Flt>,
        amplitude_modulation: Option<Vec<(Positive, Positive)>>,
    ) -> Result<Self> {
        ensure!(
            fl <= fu,
            InvalidParameterSnafu {
                param: "fl",
                criterion: "Lower frequency should be smaller than upper frequency"
            }
        );
        let quiet_time = if let Some(quiet_time) = quiet_time {
            Some(quiet_time.try_into().context(ParameterOutOfRangeSnafu {
                parameter: "Quiet time",
            })?)
        } else {
            None
        };
        Ok(SweepSettings {
            sweep_type,
            fl: fl.try_into().context(ParameterOutOfRangeSnafu {
                parameter: "Lower frequency",
            })?,
            fu: fu.try_into().context(ParameterOutOfRangeSnafu {
                parameter: "Upper frequency",
            })?,
            sweep_time: sweep_time.try_into().context(ParameterOutOfRangeSnafu {
                parameter: "Sweep time",
            })?,
            quiet_time,
            amplitude_modulation,
        })
    }

    fn getSignal(&self, fs: StrictlyPositive) -> Result<Dcol> {
        // Number of samples in sweep
        let Ns = (*self.sweep_time * *fs) as usize;
        let phase = self.getPhase(Ns, fs);
        let fs = *fs;
        // Number of samples in quiet time
        let Nq = if let Some(qt) = &self.quiet_time {
            (**qt * fs) as usize
        } else {
            0
        };

        // Total number of samples
        let N = Ns + Nq;

        Ok(Dcol::from_iter(
            (0..N).map(|i| if i < Ns { Flt::sin(phase[i]) } else { 0. }),
        ))
    }

    /// Returns the phase as a function of time
    fn getPhase(&self, Ns: usize, fs: StrictlyPositive) -> Dcol {
        match self.sweep_type {
            SweepType::BackwardLin | SweepType::ForwardLin => self.getLinSweepFBPhase(Ns, fs),
            SweepType::BackwardLog | SweepType::ForwardLog => self.getLogSweepFBPhase(Ns, fs),
            SweepType::ContinuousLin => self.getLinSweepContPhase(Ns, fs),
            SweepType::ContinuousLog => self.getLogSweepContPhase(Ns, fs),
        }
    }

    // Linear forward or backward sweep phase
    fn getLinSweepFBPhase(&self, Ns: usize, fs: StrictlyPositive) -> Dcol {
        let sweep_type = &self.sweep_type;
        assert!(matches!(
            sweep_type,
            SweepType::BackwardLin | SweepType::ForwardLin
        ));
        let (fl, fu) = (*self.fl, *self.fu);
        // For backward sweeps, we just reverse the start and stop frequency.
        let (fl, fu) = if matches!(sweep_type, SweepType::BackwardLin | SweepType::BackwardLog) {
            (fu, fl)
        } else {
            (fl, fu)
        };

        // Time step
        let Dt = 1. / *fs;
        let Nsf = Ns as Flt;
        let K = (Dt * (fl * Nsf + 0.5 * (Nsf - 1.) * (fu - fl))).floor();
        let eps_num = K / Dt - fl * Nsf - 0.5 * (Nsf - 1.) * (fu - fl);
        let eps = eps_num / (0.5 * (Nsf - 1.));
        let mut phase = 0.;
        Dcol::from_iter((0..Ns).map(|n| {
            let freq = fl + (n as Flt - 1.) / (Ns as Flt) * (fu + eps - fl);
            let phase_out = phase;
            phase += twopi * Dt * freq;
            phase_out
        }))
    }

    // Logarithmic forward or backward sweep phase
    fn getLogSweepFBPhase(&self, Ns: usize, fs: StrictlyPositive) -> Dcol {
        let sweep_type = &self.sweep_type;
        assert!(matches!(
            sweep_type,
            SweepType::BackwardLog | SweepType::ForwardLog
        ));

        let fs = *fs;
        let (fl, fu) = (*self.fl, *self.fu);
        // For backward sweeps, we just reverse the start and stop frequency.
        let (fl, fu) = if matches!(sweep_type, SweepType::BackwardLin | SweepType::BackwardLog) {
            (fu, fl)
        } else {
            (fl, fu)
        };
        // // Time step
        let Dt = 1. / fs;
        let Nsf = Ns as Flt;
        let mut k = fu / fl;
        let K = (Dt * fl * (k - 1.) / ((k.powf(1.0 / Nsf)) - 1.)).floor();

        /* Iterate k to the right solution */
        (0..10).for_each(|_| {
            let E = 1. + K / (Dt * fl) * (k.powf(1.0 / Nsf) - 1.) - k;
            let dEdk = K / (Dt * fl) * k.powf(1.0 / Nsf) / (Nsf * k) - 1.;
            k -= E / dEdk;
        });

        let mut phase = 0.;
        Dcol::from_iter((0..Ns).map(|n| {
            let nf = n as Flt;
            let fnn = fl * k.powf(nf / Nsf);
            let phase_old = phase;
            phase += twopi * Dt * fnn;
            phase_old
        }))
    }

    // Continuous log sweep phase
    fn getLogSweepContPhase(&self, Ns: usize, fs: StrictlyPositive) -> Dcol {
        let sweep_type = &self.sweep_type;
        assert!(matches!(sweep_type, SweepType::ContinuousLog));

        let fs = *fs;
        let (fl, fu) = (*self.fl, *self.fu);
        // // Time step
        let Dt = 1. / fs;
        let Nf = Ns / 2;
        let Nff = Nf as Flt;
        let Nb = Ns - Nf;
        let Nbf = Nb as Flt;
        let k1 = fu / fl;
        let phif1 = twopi * Dt * fl * (k1 - 1.) / (k1.powf(1.0 / Nff) - 1.);

        let K =
            (phif1 / twopi + Dt * fu * (1. / k1 - 1.) / ((1. / k1).powf(1.0 / Nbf) - 1.)).floor();
        let mut k = k1;

        /* Newton iterations to converge k to the value such that the sweep is
         * continuous */
        (0..NITER_NEWTON).for_each(|_| {
            let E = (k - 1.) / (k.powf(1.0 / Nff) - 1.) + (k - 1.) / (1. - k.powf(-1.0 / Nbf))
                - K / Dt / fl;

            //     /* All parts of the derivative of above error E to k */
            let dEdk1 = 1. / (k.powf(1.0 / Nff) - 1.);
            let dEdk2 = (1. / k - 1.) / (k.powf(-1.0 / Nbf) - 1.);
            let dEdk3 = -1. / (k * (k.powf(-1.0 / Nbf) - 1.));
            let dEdk4 = k.powf(-1.0 / Nbf) * (1. / k - 1.)
                / (Nbf * Flt::powi(Flt::powf(k, -1.0 / Nbf) - 1., 2));

            let dEdk5 = -Flt::powf(k, 1.0 / Nff) * (k - 1.)
                / (Nff * k * Flt::powi(Flt::powf(k, 1.0 / Nff) - 1., 2));

            let dEdk = dEdk1 + dEdk2 + dEdk3 + dEdk4 + dEdk5;
            k -= E / dEdk;
        });

        let mut phase = 0.;
        Dcol::from_iter((0..Ns).map(|n| {
            let nf = n as Flt;
            let fnn = if n <= Nf {
                fl * k.powf(nf / Nff)
            } else {
                fl * k * (1. / k).powf((nf - Nff) / Nbf)
            };
            let phase_old = phase;
            phase += twopi * Dt * fnn;

            phase_old
        }))
    }

    // Continuous linear sweep phase
    fn getLinSweepContPhase(&self, Ns: usize, fs: StrictlyPositive) -> Dcol {
        assert!(matches!(self.sweep_type, SweepType::ContinuousLin));
        let fs = *fs;
        let (fl, fu) = (*self.fl, *self.fu);
        let Dt = 1. / fs;
        let Nf = Ns / 2;
        let Nb = Ns - Nf;
        let Nff = Nf as Flt;
        let Nbf = Nb as Flt;
        /* Phi halfway */
        let phih = twopi * Dt * (fl * Nff + 0.5 * (Nff - 1.) * (fu - fl));
        let K = (phih / twopi + Dt * (fu * Nbf - (Nb as Flt - 1.) * (fu - fl))).floor();

        let eps_num1 = (K - phih / twopi) / Dt;
        let eps_num2 = -fu * Nbf + (Nbf - 1.) * (fu - fl);

        let eps = (eps_num1 + eps_num2) / (0.5 * (Nbf + 1.));
        let mut phase = 0.;
        Dcol::from_iter((0..Ns).map(|n| {
            let nf = n as Flt;
            let freq = if n < Nf {
                fl + nf / Nff * (fu - fl)
            } else {
                fu - (nf - Nff) / Nbf * (fu + eps - fl)
            };
            let phase_out = phase;
            phase += twopi * Dt * freq;
            phase_out
        }))
    }
}

#[cfg(feature = "python-bindings")]
#[cfg_attr(feature = "python-bindings", gen_stub_pymethods, pymethods)]
impl SweepSettings {
    /// See `SweepSettings::new()` for argument explanation.
    #[new]
    #[pyo3(signature=(fl, fu, sweep_time, sweep_type, quiet_time=None, amplitude_modulation=None))]
    fn py_new(
        fl: Flt,
        fu: Flt,
        sweep_time: Flt,
        sweep_type: SweepType,
        quiet_time: Option<Flt>,
        amplitude_modulation: Option<Vec<(Positive, Positive)>>,
    ) -> PyResult<Self> {
        Ok(Self::new(
            fl,
            fu,
            sweep_time,
            sweep_type,
            quiet_time,
            amplitude_modulation,
        )?)
    }
}

/// Sine sweep signal
#[derive(Debug, Clone)]
pub struct Sweep {
    settings: SweepSettings,
    fs: StrictlyPositive,
    N: usize,
    periodic_buf: Dcol,
}
impl Deref for Sweep {
    type Target = SweepSettings;
    fn deref(&self) -> &Self::Target {
        &self.settings
    }
}
impl Sweep {
    /// Create a new sine sweep signal
    ///
    /// # Args
    ///
    /// * `settings`: Sweep settings
    /// * `fs`: Sample rate
    pub fn new(settings: SweepSettings, fs: StrictlyPositive) -> Result<Self> {
        ensure!(
            1. / *settings.sweep_time < *fs / 2.,
            InvalidParameterSnafu {
                param: "Sweep time",
                criterion: "Must be less than Nyquist frequency"
            }
        );
        ensure!(*fs / 2. >= *settings.fu, SampleRateTooLowSnafu { fs: *fs });
        let periodic_buf = settings.getSignal(fs)?;
        Ok(Sweep {
            N: 0,
            settings,
            fs,
            periodic_buf,
        })
    }
}
// Linear forward or backward sweep phase
impl SourceImpl for Sweep {
    fn genSignal_unscaled(&mut self, sig: &mut dyn ExactSizeIterator<Item = &mut Flt>) {
        // New position of N is obtained here, as the iterator len() changes
        // during iteration, so we should initialize it at the start.
        let Nnew = self.N + sig.len();

        // Nice things of Rust: cycle iterators that are zipped with the output
        // iterator
        let sweep_iter = self
            .periodic_buf
            .as_slice()
            .unwrap()
            .iter()
            .cycle()
            .skip(self.N);
        for (sig, sweep_sample) in sig.zip(sweep_iter) {
            *sig = *sweep_sample;
        }

        // Update position for next start
        // Modulo number of samples in generator
        self.N = Nnew % self.periodic_buf.len();
    }

    fn fs(&self) -> StrictlyPositive {
        self.fs
    }
}

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

    use super::*;

    #[test]
    fn test_phase_linsweep1() {
        let fs = 10.;
        let fl = 1.;
        let fu = 1.;
        let sweep_settings =
            SweepSettings::new(fl, fu, 10., SweepType::ForwardLin, None, None).unwrap();

        let Ns = (*sweep_settings.sweep_time * fs) as usize;
        let phase = sweep_settings.getPhase(Ns, fs.try_into().unwrap());

        assert_abs_diff_eq!(phase[10], &(twopi));
    }
}