#[macro_use]
extern crate colour;
extern crate rand;

mod optimizer;

pub use optimizer::{job_config, swarm_config, PSO};

#[cfg(test)]
mod tests {

    use super::*;
    use job_config::JobConfig;
    use std::cell::RefCell;
    use std::f64::consts::PI;
    use swarm_config::{SwarmConfig, TransientBehavior};

    #[test]
    fn exp_sin_cost_fn() {
        // this cost function is pretty difficult to minimize quickly
        let num_variables = 3;

        // just a few iterations
        let mut jc = JobConfig::new(num_variables);
        jc.max_iterations(1000);
        jc.update_console(250);
        jc.variable_bound([-2.0 * PI, 2.0 * PI]);

        // get rekt local minima
        let mut tb = TransientBehavior::new();
        tb.stochastic_behavior(3.0, 5);
        tb.momentum_mode(0);
        tb.motion_mode(3);

        // mostly disregard global minima
        let mut sc = SwarmConfig::new();
        sc.synergic_behavior(0.4, 128);
        sc.motion_coefficients(0.5, 0.6, 1.0);
        sc.num_particles(128);
        sc.set_transient_behavior(tb);

        // lots of threads (all with the same sc)
        let pso = PSO::from_swarm_config(8, true, &sc);

        let min = pso.run_job_fn(jc, move |pt: &[f64]| -> f64 {
            let mut sum = 0.0;
            let mut sin_sum = 0.0;
            for i in 0..num_variables {
                sum += pt[i].abs();
                sin_sum += pt[i].powi(2).sin();
            }
            sum * (-1.0 * sin_sum).exp()
        });

        println!("minimum of: {}, located at: {:?}", min.0, min.1);

        assert!(min.0 < 10e-10);
    }

    #[test]
    fn mut_cost_fn() {
        let num_variables = 5;
        let num_swarms = 4;

        let mut jc = JobConfig::new(num_variables);
        jc.max_iter_and_exit_cost(100000, 10e-10);
        jc.update_console(100);
        jc.variable_bound([-10.0, 10.0]);

        let mut sc = SwarmConfig::new();
        sc.motion_coefficients(0.7, 0.9, 0.8);
        sc.num_particles(256);

        let pso = PSO::from_swarm_config(num_swarms, true, &sc);

        // !Sync, internally mutating structure for the cost function
        #[derive(Clone)]
        struct CostFuncDataThing {
            data: Vec<f64>,
            previous_cost: RefCell<f64>,
        }

        impl CostFuncDataThing {
            pub fn compute_cost(&mut self, point: &[f64]) -> f64 {
                let cost = self
                    .data
                    .iter()
                    .zip(point.iter())
                    .map(|(d, p)| (d - p).powi(2))
                    .sum::<f64>();

                // store cost for no reason at all
                self.previous_cost.replace(cost); // !sync and requires &mut self

                cost
            }
        }
        let mut cfd = CostFuncDataThing {
            data: vec![1.0; num_variables],
            previous_cost: RefCell::new(0.0),
        };

        let min = pso.run_job_fn_mut(jc, move |point: &[f64]| -> f64 { cfd.compute_cost(point) });

        println!("minimum of: {}, located at: {:?}", min.0, min.1);

        assert!(min.0 < 10e-9);
    }

    #[test]
    fn super_simple() {
        // set up a PSO with 8 default swarms
        let pso = PSO::default(8, false);

        // set up the stop condition and search space with a JobConfig
        let mut jc = JobConfig::new(5); // search a 5 variable space
        jc.exit_cost(10e-10); // stop optimizing when the objective cost reaches 10e-10
        jc.variable_bound([-5.0, 5.0]); // constrain the 5D search space to (-5, 5) along all dimensions

        // create a simple objective function
        let obj = |x_vec: &[f64]| -> f64 { x_vec.iter().map(|x| x.powi(2)).sum::<f64>() };

        // search for the minimum
        let min = pso.run_job_fn(jc, obj);

        assert!(min.0 < 10e-9);
        for x in min.1 {
            assert!(x.abs() < 0.001);
        }
    }

    #[test]
    fn advanced() {
        // define a swarm configuration
        let mut sc = SwarmConfig::new();
        sc.synergic_behavior(0.4, 100); // collaborate with other swarms every 100 iterations using a global-coefficient of 0.4
        sc.motion_coefficients(0.5, 0.6, 1.0); // use custom motion coefficients
        sc.num_particles(256); // put 256 particles in each swarm

        // set up a PSO with 8 swarms using the above SwarmConfig
        let pso = PSO::from_swarm_config(8, false, &sc);

        // set up the stop condition and search space with a JobConfig
        let mut jc = JobConfig::new(5);
        jc.max_iter_and_exit_cost(10000, 10e-10); // stop after the first of these conditions is met
        jc.variable_bounds(vec![
            // define a custom upper and lower bound on each dimension
            [-10.0, 10.0],
            [-8.0, 8.0],
            [-6.0, 6.0],
            [-4.0, 4.0],
            [-2.0, 2.0],
        ]);
        jc.update_console(100); // update the console with the current minimum every 100 iterations

        // create a simple objective function using some external data
        let mins = [1.0; 5];
        let obj = move |x_vec: &[f64]| -> f64 {
            x_vec
                .iter()
                .enumerate()
                .map(|(i, x)| (x - mins[i]).powi(2))
                .sum::<f64>()
        };

        // search for the minimum
        let min = pso.run_job_fn(jc, obj);

        assert!(min.0 < 10e-9);
        for (i, x) in min.1.iter().enumerate() {
            assert!((x - mins[i]).abs() < 0.001);
        }
    }
}