nmr-schedule 0.2.1

use core::fmt::Display;

use crate::{
    Schedule,
    generators::{Generator, Trace},
    modifier,
};

pub(crate) mod generators {
    /// The generator after applying [`super::TMFilter`]
    #[derive(Debug)]
    pub struct TMFilterGenerator<T> {
        pub(crate) generator: T,
    }
}

use alloc::vec::Vec;
use generators::*;
use ndarray::Ix1;

use super::{Filter, Modifier};

/// The trace for `TMFilter`
#[derive(Debug)]
pub struct TMFilterTrace {
    /// Gives the indices of each swap performed in the schedule
    pub swaps_done: Vec<usize>,
}

impl Display for TMFilterTrace {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "Performed {} swaps", self.swaps_done.len())
    }
}

impl<T: Generator<Ix1>> Generator<Ix1> for TMFilterGenerator<T> {
    fn _generate(&self, count: usize, dims: Ix1, iteration: u64) -> Trace<Ix1> {
        let trace = self
            .generator
            .generate_with_iter_and_trace(count, dims, iteration);

        TMFilter::filter_from_iter_and_trace(&TMFilter::new(), trace, iteration)
    }
}

modifier!(
    TMFilter<Ix1>,
    TMFilterBuilder,
    r#"Filter the schedule to reduce patterns by swapping samples to make the schedule "look like" the Thue-Morse sequence.
    
The Thue-Morse sequence is known to avoid consecutive patterns and this effectively transfers that property to the schedule.
    
The iteration parameter alters the slice of the Thue-Morse sequence that is used for filtering."#,
    tm_filter,
);

impl Modifier<Ix1> for TMFilter {
    type Output<T: Generator<Ix1>> = TMFilterGenerator<T>;

    fn modify<T: Generator<Ix1>>(self, generator: T) -> Self::Output<T> {
        TMFilterGenerator { generator }
    }
}

impl Filter<Ix1> for TMFilter {
    fn filter_with_iter_and_trace(&self, sched: Schedule<Ix1>, mut iteration: u64) -> Trace<Ix1> {
        if sched.len() < 2 {
            return Trace::new(
                sched,
                TMFilterTrace {
                    swaps_done: Vec::new(),
                },
            );
        }

        if iteration > u64::MAX - (sched.len() as u64) {
            iteration -= sched.len() as u64;
        }

        let mut sched = sched.into_inner();

        const fn thue_morse(i: u64) -> bool {
            i.count_ones() % 2 == 0
        }

        let mut swaps_done = Vec::new();

        for i in 1..sched.len() - 2 {
            let v64 = i as u64;
            if sched[i] == thue_morse(v64 + iteration + 1)
                && sched[i + 1] == thue_morse(v64 + iteration)
                && sched[i] != sched[i + 1]
            {
                swaps_done.push(i);
                sched[i] = thue_morse(v64 + iteration);
                sched[i + 1] = thue_morse(v64 + iteration + 1);
            }
        }

        Trace::new(Schedule::new(sched), TMFilterTrace { swaps_done })
    }
}

#[cfg(test)]
mod tests {
    use ndarray::Ix1;

    use crate::{
        generators::{Generator, Quantiles},
        pdf::unweighted,
    };

    use super::TMFilterBuilder;

    #[test]
    fn tm_overflow() {
        Quantiles::new(unweighted)
            .tm_filter()
            .generate_with_iter(64, Ix1(256), u64::MAX);

        Quantiles::new(unweighted)
            .tm_filter()
            .generate_with_iter(64, Ix1(256), u64::MAX - 253);
    }
}