lasprs 0.8.0

//! Sweep signal generation code
use strum::EnumMessage;
use strum_macros::{Display, EnumMessage};
use {
    crate::config::*,
    anyhow::{bail, Result},
};
const NITER_NEWTON: usize = 20;
const twopi: Flt = 2. * pi;

/// Enumerator representing the type of sweep source to create. Used as
/// parameter in [Siggen::newSweep].
#[cfg_attr(feature = "python-bindings", pyclass(eq, eq_int))]
#[derive(Debug, PartialEq, Clone, Display, EnumMessage)]
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", pymethods)]
impl SweepType {
    #[staticmethod]
    fn all() -> Vec<SweepType> {
        use SweepType::*;
        vec![
            ForwardLin,
            ForwardLog,
            BackwardLin,
            BackwardLog,
            ContinuousLin,
            ContinuousLog,
        ]
    }
    fn __str__(&self) -> String {
        self.get_message().unwrap().into()
    }
}

#[derive(Debug, Clone)]
pub struct SweepParams {
    // These parameters are described at [Source::newSweep]
    fs: Flt,
    fl: Flt,
    fu: Flt,
    sweep_time: Flt,
    quiet_time: Flt,
    sweeptype: SweepType,
}
impl SweepParams {
    pub fn new(
        fs: Flt,
        fl: Flt,
        fu: Flt,
        sweep_time: Flt,
        quiet_time: Flt,
        sweeptype: SweepType,
    ) -> Result<Self> {
        if fs <= 0. {
            bail!("Invalid sampling frequency: {} Hz", fs);
        }
        if fl > fu {
            bail!("Lower frequency should be smaller than upper frequency");
        }
        if fu >= fs / 2. {
            bail!("Upper frequency should be smaller than sampling frequency");
        }
        if sweep_time <= 0. {
            bail!("Invalid sweep time, should be > 0.");
        }
        if 1. / sweep_time > fs / 2. {
            bail!("Invalid sweep time: too short");
        }

        // For backward sweeps, we just reverse the start and stop frequency.
        let (fl, fu) = if matches!(sweeptype, SweepType::BackwardLin | SweepType::BackwardLog) {
            (fu, fl)
        } else {
            (fl, fu)
        };
        Ok(SweepParams {
            fs,
            fl,
            fu,
            sweep_time,
            quiet_time,
            sweeptype,
        })
    }
    pub fn reset(&mut self, fs: Flt) -> Dcol {
        self.fs = fs;
        self.getSignal()
    }

    fn Ns(&self) -> usize {
        (self.sweep_time * self.fs) as usize
    }
    /// Returns the phase as a function of time
    fn getPhase(&self) -> Dcol {
        match self.sweeptype {
            SweepType::BackwardLin | SweepType::ForwardLin => self.getLinSweepFBPhase(),
            SweepType::BackwardLog | SweepType::ForwardLog => self.getLogSweepFBPhase(),
            SweepType::ContinuousLin => self.getLinSweepContPhase(),
            SweepType::ContinuousLog => self.getLogSweepContPhase(),
        }
    }
    pub fn getSignal(&self) -> Dcol {
        let fs = self.fs;
        // Number of samples in sweep
        let Ns = (self.sweep_time * fs) as usize;
        // Number of samples in quiet time
        let Nq = (self.quiet_time * fs) as usize;

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

        let phase = self.getPhase();
        Dcol::from_iter((0..N).map(|i| if i < Ns { Flt::sin(phase[i]) } else { 0. }))
    }

    // Linear forward or backward sweep phase
    fn getLinSweepFBPhase(&self) -> Dcol {
        assert!(matches!(
            self.sweeptype,
            SweepType::BackwardLin | SweepType::ForwardLin
        ));
        let (Ns, fl, fu, fs) = (self.Ns(), self.fl, self.fu, self.fs);

        // 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) -> Dcol {
        assert!(matches!(
            self.sweeptype,
            SweepType::BackwardLog | SweepType::ForwardLog
        ));

        let (Ns, fl, fu, fs) = (self.Ns(), self.fl, self.fu, self.fs);
        // // 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) -> Dcol {
        assert!(matches!(self.sweeptype, SweepType::ContinuousLog));

        let (Ns, fl, fu, fs) = (self.Ns(), self.fl, self.fu, self.fs);
        // // 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) -> Dcol {
        assert!(matches!(self.sweeptype, SweepType::ContinuousLin));

        let (Ns, fl, fu, fs) = (self.Ns(), self.fl, self.fu, self.fs);
        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(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 phase = SweepParams::new(fs, fl, fu, 10., 0., SweepType::ForwardLin)
            .unwrap()
            .getLinSweepFBPhase();

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