nmr-schedule 0.2.1

use core::cmp::Ordering;
use core::fmt::Display;

use alloc::vec;
use ndarray::{Array, Ix1};
use rand::{Rng, SeedableRng};
use rand_chacha::ChaCha12Rng;

use crate::{Monotonicity, Schedule};

use super::{xor_iteration, Generator, Trace};

fn sample_poisson(λ: f64, fail_after: usize, mut rng: impl Rng) -> usize {
    let threshold = (-λ).exp();
    let mut i = 0;
    let mut u = 1.;

    loop {
        u *= rng.random::<f64>();

        if u < threshold || i > fail_after {
            return i;
        }

        i += 1;
    }
}

/// The original Sine Weighted Poisson-gap algorithm.
///
/// The iteration parameter alters the random seed.
///
/// S. G. Hyberts, K. Takeuchi, G. Wagner, J Am Chem Soc. 2010, 132(7), 2145-2147. <https://doi.org/10.1021/ja908004w>
#[derive(Clone, Copy, Debug)]
pub struct SinWeightedPoissonGap {
    seed: [u8; 32],
}

impl SinWeightedPoissonGap {
    /// Create a new `SinWeightedPoissonGap` from a random seed.
    pub const fn new(seed: [u8; 32]) -> SinWeightedPoissonGap {
        SinWeightedPoissonGap { seed }
    }
}

/// The trace for `SinWeightedPoissonGap`
#[derive(Clone, Copy, Debug)]
pub struct SinWeightedPoissonGapTrace {
    /// How many iterations were required to find a λ that gave the correct count
    pub iterations: usize,
    /// Which values of λ and `iteration` gave the correct count
    pub final_params: (f64, u64),
}

impl Display for SinWeightedPoissonGapTrace {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(
            f,
            "Generation tries needed: {}, Final parameters: ⟨λ {}, Seed XOR = {}⟩",
            self.iterations, self.final_params.0, self.final_params.1
        )
    }
}

impl Generator<Ix1> for SinWeightedPoissonGap {
    fn _generate(&self, count: usize, dims: Ix1, iteration: u64) -> Trace<Ix1> {
        if count == 0 {
            return Trace::new(
                Schedule::new(Array::from_vec(vec![false; dims[0]])),
                SinWeightedPoissonGapTrace {
                    iterations: 0,
                    final_params: (0., 0),
                },
            );
        }

        let (sched, λ, iterations) = generate_by_fuzz(
            count,
            iteration,
            Monotonicity::Decreasing,
            dims[0] as f64 / count as f64,
            |arbitrary, iteration| {
                let mut rng = ChaCha12Rng::from_seed(xor_iteration(self.seed, iteration));
                let mut schedule = alloc::vec![false; dims[0]];

                let mut i = 0;

                while i < dims[0] {
                    schedule[i] = true;

                    i += 1 + sample_poisson(
                        arbitrary
                            * ((i as f64 + 0.5) / (dims[0] as f64 + 1.)
                                * core::f64::consts::FRAC_PI_2)
                                .sin(),
                        dims[0],
                        &mut rng,
                    );
                }

                Schedule::new(Array::from_vec(schedule))
            },
        );

        Trace::new(
            sched,
            SinWeightedPoissonGapTrace {
                iterations,
                final_params: (λ),
            },
        )
    }
}

/// Generate schedules by randomly varying the `iteration` parameter while at the same time nudging an arbitrary parameter until the algorithm gives the correct number of points.
///
/// Binary search cannot be used for Poisson Gap because it does not converge for all random seeds (small variations tend to snowball and prevent convergence)
///
/// This is also the method used by the code in the supplement of the original paper.
fn generate_by_fuzz(
    count: usize,
    original_iteration: u64,
    monotonicity: Monotonicity,
    start_value: f64,
    generate_one: impl Fn(f64, u64) -> Schedule<Ix1>,
) -> (Schedule<Ix1>, (f64, u64), usize) {
    let mut rng = ChaCha12Rng::from_seed([0; 32]);
    let mut arbitrary = start_value;
    let mut iterations = 0;

    loop {
        iterations += 1;

        let iteration = if iterations == 1 {
            0
        } else {
            rng.random::<u64>()
        } ^ original_iteration;

        let sched = generate_one(arbitrary, iteration);

        let sched_count = sched.count();
        let adjust = 1. + (2.5 / (sched.len() as f64) * (count as f64 - sched_count as f64).abs());

        match (sched_count.cmp(&count), monotonicity) {
            (Ordering::Less, Monotonicity::Increasing)
            | (Ordering::Greater, Monotonicity::Decreasing) => arbitrary *= adjust,
            (Ordering::Equal, _) => return (sched, (arbitrary, iteration), iterations),
            (Ordering::Greater, Monotonicity::Increasing)
            | (Ordering::Less, Monotonicity::Decreasing) => arbitrary /= adjust,
        }
    }
}