nmr-schedule 0.2.1

use core::fmt::Display;

use alloc::vec;
use ndarray::{Array, Ix1};

use crate::{pdf::PdfGenerator, schedule::Schedule};

use super::{Generator, Trace};

/// Schedule generation by Quantiles.
///
/// L. Craft, R. Sonstrom, V. Rovnyak, D. Rovnyak, Journal of magnetic resonance, 288, 109-121. <https://doi.org/10.1016/j.jmr.2018.01.014>
#[derive(Clone, Copy, Debug)]
pub struct Quantiles<G: PdfGenerator<Ix1>>(G);

impl<G: PdfGenerator<Ix1>> Quantiles<G> {
    /// Create a new quantiles generator with the given PDF
    pub const fn new(pdf: G) -> Quantiles<G> {
        Quantiles(pdf)
    }
}

/// Trace information for `Quantiles`. Currently empty.
#[derive(Clone, Copy, Debug)]
pub struct QuantilesTrace;

impl Display for QuantilesTrace {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "No trace information")
    }
}

impl<G: PdfGenerator<Ix1>> Generator<Ix1> for Quantiles<G> {
    fn _generate_no_trace(&self, count: usize, dims: Ix1, _iteration: u64) -> Schedule<Ix1> {
        if count == 0 {
            return Schedule::new(Array::from_vec(vec![false; dims[0]]));
        }

        let pdf = self.0.get(dims).pop().unwrap();

        let integral = pdf.continuous_integral();

        let interval = 1. / (count - 1) as f64;

        let mut sched = alloc::vec![false; pdf.len()];
        let mut current_quantile = 0.;
        let mut excess = 0;

        for (spot, pos) in sched.iter_mut().enumerate() {
            // QSched is implemented with a one offset for some reason so `1.5` instead of `0.5` is what this is.
            while integral(spot as f64 + 1.5) >= current_quantile * interval {
                if *pos {
                    excess += 1;
                }

                *pos = true;
                current_quantile += 1.;
            }
        }

        let last = sched.len() - 1;
        sched[last] = true;

        let mut pos = 0;

        while excess > 0 {
            if !sched[pos] {
                sched[pos] = true;
                excess -= 1;
            }

            pos += 1;
        }

        Schedule::new(Array::from_vec(sched))
    }

    fn _generate(&self, count: usize, dims: Ix1, iteration: u64) -> Trace<Ix1> {
        Trace::new(
            self._generate_no_trace(count, dims, iteration),
            QuantilesTrace,
        )
    }
}

#[cfg(test)]
mod tests {
    use core::f64::consts::PI;

    use ndarray::Ix1;

    use crate::{
        generators::Generator,
        modifiers::FillCornersBuilder,
        pdf::{exponential, qsin, QSinBias},
        EncodingType, Schedule,
    };

    use super::Quantiles;

    #[test]
    fn qsched_compatibility() {
        let schedules = [
            (
                include_str!("./tests/qsched_compat/32x128-low.sch"),
                Quantiles::new(|len| qsin(len, QSinBias::Low, PI)).generate(32, Ix1(128)),
            ),
            (
                include_str!("./tests/qsched_compat/32x128-med.sch"),
                Quantiles::new(|len| qsin(len, QSinBias::Med, 3.)).generate(32, Ix1(128)),
            ),
            (
                include_str!("./tests/qsched_compat/32x128-high.sch"),
                Quantiles::new(|len| qsin(len, QSinBias::High, 3.)).generate(32, Ix1(128)),
            ),
            (
                include_str!("./tests/qsched_compat/32x128.sch"),
                Quantiles::new(|len| qsin(len, QSinBias::Med, 3.))
                    .fill_corners(|_, _| [8, 1])
                    .generate(32, Ix1(128)),
            ),
            (
                include_str!("./tests/qsched_compat/42x128.sch"),
                Quantiles::new(|len| exponential(len, 3.))
                    .fill_corners(|_, _| [10, 1])
                    .generate(42, Ix1(128)),
            ),
            (
                include_str!("./tests/qsched_compat/64x256.sch"),
                Quantiles::new(|len| qsin(len, QSinBias::Med, 3.))
                    .fill_corners(|_, _| [12, 0])
                    .generate(64, Ix1(256)),
            ),
            (
                include_str!("./tests/qsched_compat/85x256.sch"),
                Quantiles::new(|len| exponential(len, 3.))
                    .fill_corners(|_, _| [18, 1])
                    .generate(85, Ix1(256)),
            ),
            (
                include_str!("./tests/qsched_compat/130x512.sch"),
                Quantiles::new(|len| qsin(len, QSinBias::Med, 3.))
                    .fill_corners(|_, _| [15, 0])
                    .generate(130, Ix1(512)),
            ),
            (
                include_str!("./tests/qsched_compat/169x512.sch"),
                Quantiles::new(|len| exponential(len, 3.))
                    .fill_corners(|_, _| [15, 1])
                    .generate(169, Ix1(512)),
            ),
            (
                include_str!("./tests/qsched_compat/257x1024.sch"),
                Quantiles::new(|len| qsin(len, QSinBias::Med, 3.))
                    .fill_corners(|_, _| [30, 0])
                    .generate(257, Ix1(1024)),
            ),
            (
                include_str!("./tests/qsched_compat/337x1024.sch"),
                Quantiles::new(|len| exponential(len, 3.))
                    .fill_corners(|_, _| [30, 1])
                    .generate(337, Ix1(1024)),
            ),
        ];

        for (truth, generated) in schedules.into_iter() {
            let truth = Schedule::decode(truth, EncodingType::ZeroBased, |v: Ix1| Ok(v)).unwrap();

            assert_eq!(truth.count(), generated.count());
            assert_eq!(truth, generated);
        }
    }
}