Mercurial > repos > alg_tools / file revision

/*!

This module provides tools for separating the computational steps of iterative
algorithms from visualisation and other reporting routines. The approach is heavily
indebted to functional programming. The computational step is to be implemented as a
closure. That closure gets passed a `state` parameter that implements an
[`if_verbose`][AlgIteratorState::if_verbose] method that is to be called to determine
whether function values or other potentially costly things need to be calculated on that
iteration. The parameter of [`if_verbose`][AlgIteratorState::if_verbose] is another
closure that does the necessary computation.

## Simple example

```rust
use alg_tools::types::*;
use alg_tools::iterate::*;
let mut iter = AlgIteratorOptions{ max_iter : 100,
                                   verbose_iter : Verbose::Every(10),
                                   .. Default::default() };
let mut x = 1 as float;
iter.iterate(|state|{
    // This is our computational step
    x = x + x.sqrt();
    state.if_verbose(||{
        // return current value when requested
        return x
     })
})
```
There is no colon after `state.if_verbose`, because we need to return its value. If you do something after the step, you need to store the result in a variable and return it from the main computational step afterwards.

This will print out
```output
10/100 J=31.051164
20/100 J=108.493699
30/100 J=234.690039
40/100 J=410.056327
50/100 J=634.799262
60/100 J=909.042928
70/100 J=1232.870172
80/100 J=1606.340254
90/100 J=2029.497673
100/100 J=2502.377071
```
*/

use colored::{Colorize,ColoredString};
//use num_traits::Num;
use core::fmt::Debug;
//use core::fmt::Display;
//use std::sync::mpsc::Receiver;
use serde::{Serialize, Deserialize};
use cpu_time::ProcessTime;
use std::marker::PhantomData;
use std::time::Duration;
use crate::types::*;
use crate::logger::*;

/// Create the displayed presentation for log items. Override for nice (condensed) presentation of rich log items, or define `show` if it is not defined for your type.
pub trait LogRepr : Debug {
    fn logrepr(&self) -> ColoredString { format!("« {:?} »", self).as_str().into() }
}

impl LogRepr for str {
    fn logrepr(&self) -> ColoredString { self.into() }
}

impl LogRepr for String {
    fn logrepr(&self) -> ColoredString { self.as_str().into() }
}

impl<T> LogRepr for T where T : Num {
    fn logrepr(&self) -> ColoredString { format!("J={}", self).as_str().into() }
}

impl<V> LogRepr for Option<V> where V : LogRepr {
    fn logrepr(&self) -> ColoredString {
        match self {
            None    => { "===missing value===".red() }
            Some(v) => { v.logrepr() }
        }
    }
}

#[derive(Debug,Clone)]
pub struct Annotated<F>(pub F, pub String);

impl<V> LogRepr for Annotated<V> where V : LogRepr {
    fn logrepr(&self) -> ColoredString {
        format!("{}\t| {}", self.0.logrepr(), self.1).as_str().into()
    }
}


/// Basic log item.
#[derive(Serialize, Deserialize, Debug, Eq, PartialEq, Clone)]
pub struct LogItem<V> {
    pub iter : usize,
    // This causes [`csv`] to crash.
    //#[serde(flatten)]
    pub data : V
}

fn make_logitem<V>(iter : usize, data : V) -> LogItem<V> {
    LogItem{ iter, data }
}

/// State of an [`AlgIterator`]
pub trait AlgIteratorState {
    fn if_verbose<A>(&self, calc_objective : impl FnMut() -> A) -> Option<A>;

    fn iteration(&self) -> usize;
}

/// Result of a step of an [`AlgIterator`]
#[derive(Debug, Serialize)]
pub enum Step<V> {
    /// Iteration should be terminated
    Terminated,
    /// No result this iteration (due to verbosity settings)
    Quiet,
    /// Result of this iteration (due to verbosity settings)
    Result(V),
}

/// An iterator for algorithms, produced by [`AlgIteratorFactory.prepare`].
/// For usage instructions see the help for [`AlgIteratorFactory`].
pub trait AlgIterator<V> : Sized {
    type State : AlgIteratorState;
    type Output;

    /// Advance the iterator.
    fn step(&mut self) -> Step<Self::Output>;

    /// Return current iteration count.
    fn iteration(&self) -> usize {
        self.state().iteration()
    }

    /// Return current iteration stats.
    fn state(&self) -> Self::State;

    /// Iterate the function `step`. It should accept one `state` ([`AlgIteratorState`] parameter,
    /// and return the output of `state.if_verbose`.
    /// Typically called through [`AlgIteratorFactory::iterate`].
    #[inline]
    fn iterate(&mut self) {
        loop {
            if let Step::Terminated = self.step() {
                break
            }
        }
    }

    /// Converts the `AlgIterator` into a plain [`Iterator`], discarding [`Step::Quiet`] results.
    /// (This is to avoid having to manually implement [`IntoIterator`] or [`Iterator`]
    /// for every single [`AlgIterator`].)
    fn downcast(self) -> AlgIteratorI<Self, V> {
        AlgIteratorI(self, PhantomData)
    }
}

/// Conversion of an `AlgIterator` into a plain [`Iterator`].
/// It discards [`Step::Quiet`] results.
pub struct AlgIteratorI<A, V>(A, PhantomData<V>);

impl<A, V> Iterator for AlgIteratorI<A, V>
where A : AlgIterator<V> {
    type Item = A::Output;

    fn next(&mut self) -> Option<A::Output> {
        loop {
            match self.0.step() {
                Step::Result(v) => return Some(v),
                Step::Terminated => return None,
                Step::Quiet => continue,
            }
        }
    }
}

/// A factory for producing an [`AlgIterator`]. To use one:
///
/// ```rust
/// use alg_tools::iterate::*;
/// let mut iter = AlgIteratorOptions::default();
/// iter.iterate(|state|{
///     // perform iterations
///     state.if_verbose(||{
///         // calculate and return function value or other displayed data v
///         return 0
///     })
/// })
/// ```
pub trait AlgIteratorFactory<V> : Sized {
    type Iter<F> : AlgIterator<V, State = Self::State, Output = Self::Output>
        where F : FnMut(&Self::State) -> Step<V>;
    type State : AlgIteratorState;
    type Output;

    /// Prepare an [`AlgIterator`], consuming the factory.
    /// The function `step_fn` should accept a `state` ([`AlgIteratorState`] parameter, and return
    /// a [`Step`].
    fn prepare<F>(self, step_fn : F) -> Self::Iter<F>
    where F : FnMut(&Self::State) -> Step<V>;

    /// Iterate the function `step`. It should accept one `state` ([`AlgIteratorState`] parameter,
    /// and return the output of `state.if_verbose`. For usage instructions see the
    /// [module documentation][self].
    ///
    /// This method is equivalent to [`Self::prepare`] followed by [`AlgIterator.iterate`].
    #[inline]
    fn iterate<F>(self, mut step : F) -> () where F : FnMut(&Self::State) -> Option<V> {
        self.prepare(move |state| {
            match step(state) {
                None => Step::Quiet,
                Some(v) => Step::Result(v),
            }
        }).iterate()
    }

    /// Iterate the function `step`. It should accept a `state` ([`AlgIteratorState`] parameter,
    /// and a data parameter taking frmo `datasource`, and return the output of `state.if_verbose`.
    /// For usage instructions see the [module documentation][self].
    #[inline]
    fn iterate_data<F, D, I>(self, mut datasource : I, mut step : F)
    where F : FnMut(&Self::State, D) -> Option<V>,
          I : Iterator<Item = D> {
        self.prepare(move |state| {
            match datasource.next() {
                None => Step::Terminated,
                Some(d) => match step(state, d) {
                    None => Step::Quiet,
                    Some(v) => Step::Result(v),
                }
            }
        }).iterate()
    }

    // fn make_iterate<'own>(self)
    // -> Box<dyn (FnMut(dyn FnMut(&Self::State) -> Option<V>) -> ()) + 'own> {
    //     Box::new(move |step| self.iterate(step))
    // }

    // fn make_iterate_data<'own, D, I>(self, datasource : I)
    // -> Box<dyn (FnMut(dyn FnMut(&Self::State, D) -> Option<V>) -> ()) + 'own>
    // where I : Iterator<Item = D> + 'own {
    //     Box::new(move |step| self.iterate_data(step, datasource))
    // }

    /// Add logging to the iterator produced by the factory.
    fn into_log<'log>(self, logger : &'log mut Logger<Self::Output>)
    -> LoggingIteratorFactory<'log, Self::Output, Self>
    where Self : Sized {
        LoggingIteratorFactory {
            base_options : self,
            logger,
        }
    }

    /// Map the output of the iterator produced by the factory.
    fn mapped<U, G>(self, map : G)
    -> MappingIteratorFactory<G, Self>
    where Self : Sized,
          G : Fn(usize, Self::Output) -> U {
        MappingIteratorFactory {
            base_options : self,
            map
        }
    }

    /// Adds iteration number to the output. Typically followed by [`Self::into_log`].
    fn with_iteration_number(self)
    -> MappingIteratorFactory<fn(usize, Self::Output) -> LogItem<Self::Output>, Self>
    where Self : Sized {
        self.mapped(make_logitem)
    }

    /// Add timing to the iterator produced by the factory.
    fn timed(self) -> TimingIteratorFactory<Self>
    where Self : Sized {
        TimingIteratorFactory(self)
    }

    /// Add value stopping threshold to the iterator produce by the factory
    fn stop_target(self, target : Self::Output) -> ValueIteratorFactory<Self::Output, Self>
    where Self : Sized,
          Self::Output : Num {
        ValueIteratorFactory { base_options : self, target : target }
    }

    /// Add stall stopping to the iterator produce by the factory
    fn stop_stall(self, stall : Self::Output) -> StallIteratorFactory<Self::Output, Self>
    where Self : Sized,
          Self::Output : Num {
        StallIteratorFactory { base_options : self, stall : stall }
    }

    /// Is the iterator quiet, i.e., on-verbose?
    fn is_quiet(&self) -> bool { false }
}

/// Options for [`BasicAlgIteratorFactory`]. Use as:
/// ```
/// # use alg_tools::iterate::*;
/// let iter = AlgIteratorOptions{ max_iter : 10000, .. Default::default() };
/// ```
#[derive(Clone, Copy, Debug, Serialize, Deserialize, Eq, PartialEq)]
pub struct AlgIteratorOptions {
    /// Maximum number of iterations
    pub max_iter : usize,
    /// Number of iterations between verbose iterations that display state.
    pub verbose_iter : Verbose,
    /// Whether verbose iterations are displayed, or just passed onwards to a containing
    /// `AlgIterator`.
    pub quiet : bool,
}

#[derive(Clone, Copy, Debug, Serialize, Deserialize, Eq, PartialEq)]
pub enum Verbose {
    /// Be verbose every $n$ iterations.
    Every(usize),
    /// Be verbose every $n$ iterations and initial $m$ iterations.
    EveryAndInitial{ every : usize, initial : usize },
    /// Be verbose if iteration number $n$ divides by $b^{\text{floor}(\log_b(n))}$, where
    /// $b$ is indicated logarithmic base. So, with $b=10$,
    /// * every iteration for first 10 iterations,
    /// * every 10 iterations from there on until 100 iterations,
    /// * every 100 iteartinos frmo there on until 1000 iterations, etc.
    Logarithmic(usize),
}

impl Verbose {
    /// Indicates whether given iteration number is verbose
    pub fn is_verbose(&self, iter : usize) -> bool {
        match self {
            &Verbose::Every(every) => {
                every != 0 && iter % every == 0
            },
            &Verbose::EveryAndInitial{ every, initial } => {
                iter <= initial || (every != 0 && iter % every == 0)
            },
            &Verbose::Logarithmic(base) => {
                let every = 10usize.pow((iter as float).log(base as float).floor() as u32);
                iter % every == 0
            }
        }
    }
}

impl Default for AlgIteratorOptions {
    fn default() -> AlgIteratorOptions {
        AlgIteratorOptions{
            max_iter : 1000,
            verbose_iter : Verbose::EveryAndInitial { every : 100, initial : 10 },
            quiet : false
        }
    }
}

/// State of a `BasicAlgIterator`
#[derive(Clone,Copy,Debug,Serialize,Eq,PartialEq)]
pub struct BasicState {
    iter : usize,
    verbose : bool,
    calc : bool,
}

/// [`AlgIteratorFactory`] for [`BasicAlgIterator`]
#[derive(Clone,Debug)]
pub struct BasicAlgIteratorFactory<V> {
    options : AlgIteratorOptions,
    phantoms : PhantomData<V>,
}

/// The simplest [`AlgIterator`], reated by [`BasicAlgIteratorFactory`]
#[derive(Clone,Debug)]
pub struct BasicAlgIterator<F> {
    options : AlgIteratorOptions,
    iter : usize,
    step_fn : F,
}

impl AlgIteratorOptions {
    /// [`AlgIteratorOptions`] is directly a factory for [`BasicAlgIterator`],
    /// however, due to type inference issues, it may become convenient to instantiate
    /// it to a specific return type for the inner step function. This method does that.
    pub fn instantiate<V>(&self) -> BasicAlgIteratorFactory<V> {
        BasicAlgIteratorFactory {
            options : self.clone(),
            phantoms : PhantomData
        }
    }
}

impl<V> AlgIteratorFactory<V> for AlgIteratorOptions
where V : LogRepr {
    type State = BasicState;
    type Iter<F> = BasicAlgIterator<F> where F : FnMut(&Self::State) -> Step<V>;
    type Output = V;

    fn prepare<F>(self, step_fn : F) -> Self::Iter<F>
    where F : FnMut(&Self::State) -> Step<V> {
        BasicAlgIterator{ options : self, iter : 0, step_fn }
    }

    #[inline]
    fn is_quiet(&self) -> bool {
        self.quiet
    }
}

impl<V> AlgIteratorFactory<V> for BasicAlgIteratorFactory<V>
where V : LogRepr {
    type State = BasicState;
    type Iter<F> = BasicAlgIterator<F> where F : FnMut(&Self::State) -> Step<V>;
    type Output = V;

    fn prepare<F>(self, step_fn : F) -> Self::Iter<F>
    where F : FnMut(&Self::State) -> Step<V> {
        BasicAlgIterator{ options : self.options, iter : 0, step_fn }
    }

    #[inline]
    fn is_quiet(&self) -> bool {
        self.options.quiet
    }
}

impl<V, F> AlgIterator<V> for BasicAlgIterator<F>
where V : LogRepr,
      F : FnMut(&BasicState) -> Step<V> {
    type State = BasicState;
    type Output = V;

    #[inline]
    fn step(&mut self) -> Step<V> {
        if self.iter >= self.options.max_iter {
            Step::Terminated
        } else {
            self.iter += 1;
            let state = self.state();
            let res = (self.step_fn)(&state);
            if let Step::Result(ref val) = res {
                if state.verbose && !self.options.quiet {
                    println!("{}{}/{} {}{}", "".dimmed(),
                            state.iter,
                            self.options.max_iter,
                            val.logrepr(),
                            "".clear());
                }
            }
            res
        }
    }

    #[inline]
    fn iteration(&self) -> usize {
        self.iter
    }

    #[inline]
    fn state(&self) -> BasicState {
        let iter = self.iter;
        let verbose = self.options.verbose_iter.is_verbose(iter);
        BasicState{
            iter : iter,
            verbose : verbose,
            calc : verbose,
        }
    }
}

impl AlgIteratorState for BasicState {
    #[inline]
    fn if_verbose<A>(&self, mut calc_objective : impl FnMut() -> A) -> Option<A> {
        if self.calc {
            Some(calc_objective())
        } else {
            None
        }
    }

    #[inline]
    fn iteration(&self) -> usize {
        return self.iter;
    }
}

//
// Stall detecting iteration function.
//

/// An [`AlgIteratorFactory`] for an [`AlgIterator`] that detects “stall”
/// $(v_{k+n}-v_k)/v_k ≤ θ$, where $n$ the distance between [`Step::Result`] iterations.
#[derive(Clone,Copy,Debug,Serialize,Eq,PartialEq)]
pub struct StallIteratorFactory<U : Num, BaseFactory> {
    /// Basic options
    pub base_options : BaseFactory,
    /// Stalling threshold $θ$.
    pub stall : U,
}

/// Iterator produced by [`StallIteratorFactory`].
pub struct StallIterator<U : Num, BaseIterator> {
    base_iterator : BaseIterator,
    stall : U,
    previous_value : Option<U>,
}

impl<V, U, BaseFactory> AlgIteratorFactory<V>
for StallIteratorFactory<U, BaseFactory>
where BaseFactory : AlgIteratorFactory<V, Output=U>,
      U : SignedNum + PartialOrd {
    type Iter<F> = StallIterator<U, BaseFactory::Iter<F>>
        where F : FnMut(&Self::State) -> Step<V>;
    type State = BaseFactory::State;
    type Output = BaseFactory::Output;

    fn prepare<F>(self, step_fn : F) -> Self::Iter<F>
    where F : FnMut(&Self::State) -> Step<V> {
        StallIterator {
            base_iterator : self.base_options.prepare(step_fn),
            stall : self.stall,
            previous_value : None,
        }
    }

    fn is_quiet(&self) -> bool {
        self.base_options.is_quiet()
    }
}

impl<V, U, BaseIterator> AlgIterator<V>
for StallIterator<U, BaseIterator>
where BaseIterator : AlgIterator<V, Output=U>,
      U : SignedNum + PartialOrd {
    type State = BaseIterator::State;
    type Output = BaseIterator::Output;

    #[inline]
    fn step(&mut self) -> Step<U> {
        match self.base_iterator.step() {
            Step::Result(nv) => {
                let previous_v = self.previous_value;
                self.previous_value = Some(nv);
                match previous_v {
                    Some(pv) if (nv - pv).abs() <= self.stall * pv.abs() => Step::Terminated,
                    _ => Step::Result(nv),
                }
            },
            val => val,
        }
    }

    #[inline]
    fn iteration(&self) -> usize {
        self.base_iterator.iteration()
    }

    #[inline]
    fn state(&self) -> Self::State {
        self.base_iterator.state()
    }
}

/// An [`AlgIteratorFactory`] for an [`AlgIterator`] that detect whether step function
/// return value is less than `target`, and terminates if it is.
#[derive(Clone,Copy,Debug,Serialize,Eq,PartialEq)]
pub struct ValueIteratorFactory<U : Num, BaseFactory> {
    /// Basic options
    pub base_options : BaseFactory,
    /// Stalling threshold $θ$.
    pub target : U,
}

/// Iterator produced by [`ValueIteratorFactory`].
pub struct ValueIterator<U : Num, BaseIterator> {
    base_iterator : BaseIterator,
    target : U,
}

impl<V, U, BaseFactory> AlgIteratorFactory<V>
for ValueIteratorFactory<U, BaseFactory>
where BaseFactory : AlgIteratorFactory<V, Output=U>,
      U : SignedNum + PartialOrd {
    type Iter<F> = ValueIterator<U, BaseFactory::Iter<F>>
        where F : FnMut(&Self::State) -> Step<V>;
    type State = BaseFactory::State;
    type Output = BaseFactory::Output;

    fn prepare<F>(self, step_fn : F) -> Self::Iter<F>
    where F : FnMut(&Self::State) -> Step<V> {
        ValueIterator {
            base_iterator : self.base_options.prepare(step_fn),
            target : self.target
        }
    }

    fn is_quiet(&self) -> bool {
        self.base_options.is_quiet()
    }
}

impl<V, U, BaseIterator> AlgIterator<V>
for ValueIterator<U, BaseIterator>
where BaseIterator : AlgIterator<V, Output=U>,
      U : SignedNum + PartialOrd {
    type State = BaseIterator::State;
    type Output = BaseIterator::Output;

    #[inline]
    fn step(&mut self) -> Step<U> {
        match self.base_iterator.step() {
            Step::Result(v) => {
                 if v <= self.target {
                    Step::Terminated
                 } else {
                    Step::Result(v)
                 }
            },
            val => val,
        }
    }

    #[inline]
    fn iteration(&self) -> usize {
        self.base_iterator.iteration()
    }

    #[inline]
    fn state(&self) -> Self::State {
        self.base_iterator.state()
    }
}

//
// Logging iterator
//

/// [`AlgIteratorFactory`] for a logging [`AlgIterator`].
/// The `Output` of the corresponding [`LoggingIterator`] is `()`:
/// The data is only inserted into the log, and not passed onwards.
/// Use as:
/// ```rust
/// use alg_tools::iterate::*;
/// use alg_tools::logger::*;
/// let iter = AlgIteratorOptions::default();
/// let mut log = Logger::new();
/// iter.into_log(&mut log).iterate(|state|{
///     // perform iterations
///     state.if_verbose(||{
///         // calculate and return function value or other displayed data v
///         return 0
///     })
/// })
/// ```
#[derive(Debug)]
pub struct LoggingIteratorFactory<'log, U, BaseFactory> {
    pub base_options : BaseFactory,
    pub logger : &'log mut Logger<U>,
}

/// Iterator produced by `LoggingIteratorFactory`.
pub struct LoggingIterator<'log, U, BaseIterator> {
    base_iterator : BaseIterator,
    logger : &'log mut Logger<U>,
}


impl<'log, V, BaseFactory> AlgIteratorFactory<V>
for LoggingIteratorFactory<'log, BaseFactory::Output, BaseFactory>
where BaseFactory : AlgIteratorFactory<V>,
      BaseFactory::Output : 'log {
    type State = BaseFactory::State;
    type Iter<F> = LoggingIterator<'log, BaseFactory::Output, BaseFactory::Iter<F>>
        where F : FnMut(&Self::State) -> Step<V>;
    type Output = ();

    fn prepare<F>(self, step_fn : F) -> Self::Iter<F>
    where F : FnMut(&Self::State) -> Step<V> {
        LoggingIterator {
            base_iterator : self.base_options.prepare(step_fn),
            logger : self.logger,
        }
    }

    #[inline]
    fn is_quiet(&self) -> bool {
        self.base_options.is_quiet()
    }
}

impl<'log, V, BaseIterator> AlgIterator<V>
for LoggingIterator<'log, BaseIterator::Output, BaseIterator>
where BaseIterator : AlgIterator<V>,
      BaseIterator::Output : 'log {
    type State = BaseIterator::State;
    type Output = ();

    #[inline]
    fn step(&mut self) -> Step<Self::Output> {
        match self.base_iterator.step() {
            Step::Result(v) => {
                self.logger.log(v);
                Step::Quiet
            },
            Step::Quiet => Step::Quiet,
            Step::Terminated => Step::Terminated,
        }
    }

    #[inline]
    fn iteration(&self) -> usize {
        self.base_iterator.iteration()
    }

    #[inline]
    fn state(&self) -> Self::State {
        self.base_iterator.state()
    }
}

/// This [`AlgIteratorFactory`] allows output mapping.
#[derive(Debug)]
pub struct MappingIteratorFactory<G, BaseFactory> {
    pub base_options : BaseFactory,
    pub map : G,
}

/// Iterator produced by `MappingIteratorFactory`.
pub struct MappingIterator<G, BaseIterator> {
    base_iterator : BaseIterator,
    map : G,
}


impl<V, U, G, BaseFactory> AlgIteratorFactory<V>
for MappingIteratorFactory<G, BaseFactory>
where BaseFactory : AlgIteratorFactory<V>,
      G : Fn(usize, BaseFactory::Output) -> U {
    type State = BaseFactory::State;
    type Iter<F> = MappingIterator<G, BaseFactory::Iter<F>>
        where F : FnMut(&Self::State) -> Step<V>;
    type Output = U;

    fn prepare<F>(self, step_fn : F) -> Self::Iter<F>
    where F : FnMut(&Self::State) -> Step<V> {
        MappingIterator {
            base_iterator : self.base_options.prepare(step_fn),
            map : self.map
        }
    }

    #[inline]
    fn is_quiet(&self) -> bool {
        self.base_options.is_quiet()
    }
}

impl<V, U, G, BaseIterator> AlgIterator<V>
for MappingIterator<G, BaseIterator>
where BaseIterator : AlgIterator<V>,
      G : Fn(usize, BaseIterator::Output) -> U {
    type State = BaseIterator::State;
    type Output = U;

    #[inline]
    fn step(&mut self) -> Step<Self::Output> {
        match self.base_iterator.step() {
            Step::Result(v) => Step::Result((self.map)(self.iteration(), v)),
            Step::Quiet => Step::Quiet,
            Step::Terminated => Step::Terminated,
        }
    }

    #[inline]
    fn iteration(&self) -> usize {
        self.base_iterator.iteration()
    }

    #[inline]
    fn state(&self) -> Self::State {
        self.base_iterator.state()
    }
}

//
// Timing iterator
//

/// An [`AlgIteratorFactory`] for an [`AlgIterator`] that adds CPU time spent to verbose events.
#[derive(Debug)]
pub struct TimingIteratorFactory<BaseFactory>(pub BaseFactory);

/// Iterator produced by [`TimingIteratorFactory`]
#[derive(Debug)]
pub struct TimingIterator<BaseIterator> {
    base_iterator : BaseIterator,
    start_time : ProcessTime,
}

/// Data `U` with production time attached
#[derive(Copy, Clone, Debug, Serialize)]
pub struct Timed<U> {
    pub cpu_time : Duration,
    //#[serde(flatten)]
    pub data : U
}

impl<T> LogRepr for Timed<T> where T : LogRepr {
    fn logrepr(&self) -> ColoredString {
        format!("[{:.3}s] {}", self.cpu_time.as_secs_f64(), self.data.logrepr()).as_str().into()
    }
}


impl<V, BaseFactory> AlgIteratorFactory<V>
for TimingIteratorFactory<BaseFactory>
where BaseFactory : AlgIteratorFactory<V> {
    type State = BaseFactory::State;
    type Iter<F> = TimingIterator<BaseFactory::Iter<F>>
        where F : FnMut(&Self::State) -> Step<V>;
    type Output = Timed<BaseFactory::Output>;

    fn prepare<F>(self, step_fn : F) -> Self::Iter<F>
    where F : FnMut(&Self::State) -> Step<V> {
        TimingIterator {
            base_iterator : self.0.prepare(step_fn),
            start_time : ProcessTime::now()
        }
    }

    #[inline]
    fn is_quiet(&self) -> bool {
        self.0.is_quiet()
    }
}

impl<V, BaseIterator> AlgIterator<V>
for TimingIterator<BaseIterator>
where BaseIterator : AlgIterator<V> {
    type State = BaseIterator::State;
    type Output = Timed<BaseIterator::Output>;

    #[inline]
    fn step(&mut self) -> Step<Self::Output> {
        match self.base_iterator.step() {
            Step::Result(data) => {
                Step::Result(Timed{
                    cpu_time : self.start_time.elapsed(),
                    data
                })
            },
            Step::Quiet => Step::Quiet,
            Step::Terminated => Step::Terminated,
        }
    }

    #[inline]
    fn iteration(&self) -> usize {
        self.base_iterator.iteration()
    }

    #[inline]
    fn state(&self) -> Self::State {
        self.base_iterator.state()
    }
}

//
// Tests
//

#[cfg(test)]
mod tests {
    use super::*;
    use crate::logger::Logger;
    #[test]
    fn iteration() {
        let options = AlgIteratorOptions{
            max_iter : 10,
            verbose_iter : Verbose::Every(3),
            .. Default::default()
        };

        {
            let mut start = 1 as int;
            options.iterate(|state| {
                start = start * 2;
                state.if_verbose(|| start)
            });
            assert_eq!(start, (2 as int).pow(10));
        }

        {
            let mut start = 1 as int;
            let mut log = Logger::new();
            let factory = options.instantiate()
                                 .with_iteration_number()
                                 .into_log(&mut log);
            factory.iterate(|state| {
                start = start * 2;
                state.if_verbose(|| start)
            });
            assert_eq!(start, (2 as int).pow(10));
            assert_eq!(log.data()
                          .iter()
                          .map(|LogItem{ data : v, iter : _ }| v.clone())
                          .collect::<Vec<int>>(),
                      (1..10).map(|i| (2 as int).pow(i))
                             .skip(2)
                             .step_by(3)
                             .collect::<Vec<int>>())
        }
    }
}