pso 0.2.0

use rand::prelude::ThreadRng;
use rand::Rng;
use std::fmt;

const DEFAULT_WALL_BOUNCE: f64 = 0.125;
const DEFAULT_LOCAL: f64 = 0.8;
const DEFAULT_TRIBAL: f64 = 1.2;
const DEFAULT_GLOBAL: f64 = 0.9;
const DEFAULT_COLLAB_RATE: usize = 16;
const DEFAULT_MOMENTUM: f64 = 0.5;
const DEFAULT_NUM_PARTICLES: usize = 128;

/// Configuration for Swarms of Particles
///
/// * Defines the physical behavior of particles
/// * Defines the number of particles in each swarm
/// * Defines how a swarm uses information from other swarms
///
/// # Examples
/// ```
/// use pso::swarm_config::SwarmConfig;
///
/// // Start with a default configuration
/// let mut sc = SwarmConfig::new();
///
/// // Tell the swarm to collaborate with others every 25 iterations and weight the global minimum with a coefficient of 0.6
/// sc.synergic_behavior(0.6, 25);
///
/// // Define other motion coefficients
/// sc.motion_coefficients(0.7, 0.9, 0.45);
///
/// // Set the number of particles in the swarm
/// sc.num_particles(256);
///
/// println!("{}", sc);
/// ```
#[derive(Clone, Debug)]
pub struct SwarmConfig {
    pub global_behavior: GlobalBehavior,
    pub transient_behavior: TransientBehavior,
    pub wall_bounce_factor: f64,
    pub local_motion: f64,
    pub tribal_motion: f64,
    pub momentum: f64,
    pub num_particles: usize,
}

impl SwarmConfig {
    pub fn new() -> Self {
        Self {
            global_behavior: GlobalBehavior::default(),
            transient_behavior: TransientBehavior::new(),
            wall_bounce_factor: DEFAULT_WALL_BOUNCE,
            local_motion: DEFAULT_LOCAL,
            tribal_motion: DEFAULT_TRIBAL,
            momentum: DEFAULT_MOMENTUM,
            num_particles: DEFAULT_NUM_PARTICLES,
        }
    }

    /// Set the particles in the swarm to use the global minimum to calculate their velocity each iteration
    ///
    /// # Arguments
    /// * `global_motion_coefficient`: How heavily weighted is the global minimum in the velocity calculation
    /// * `collaboration_periodicity`: How often does the swarm update its global minimum from the other swarms
    pub fn synergic_behavior(
        &mut self,
        global_motion_coefficient: f64,
        collaboration_periodicity: usize,
    ) {
        self.global_behavior =
            GlobalBehavior::Synergic(global_motion_coefficient, collaboration_periodicity);
    }

    /// Set the swarm to ignore the global minimum and only use its own internal minimum (this is default behavior)
    pub fn solitary_behavior(&mut self) {
        self.global_behavior = GlobalBehavior::Solitary;
    }

    /// Apply a custom transient behavior to the swarm configuration
    pub fn set_transient_behavior(&mut self, transient_behavior: TransientBehavior) {
        self.transient_behavior = transient_behavior;
    }

    /// Set the wall bounce factor. This value is multiplied by a particles velocity when it goes out of bounds.
    pub fn wall_bounce(&mut self, wall_bounce_factor: f64) {
        assert!(
            wall_bounce_factor > 0.0,
            "Wall Bounce Factor must be greater than zero!"
        );
        self.wall_bounce_factor = wall_bounce_factor * -1.0;
    }

    /// Set the motion coefficients. These describe how heavily certain parts of each particles velocity are weighted
    ///
    /// # Arguments
    /// * `local_motion`: How important is the particles best known location in the search space
    /// * `tribal_motion`: How important is the swarms best known location in the search space
    /// * `momentum`: How much of the particles current velocity is used to calculate its next velocity
    pub fn motion_coefficients(&mut self, local_motion: f64, tribal_motion: f64, momentum: f64) {
        self.local_motion = local_motion;
        self.tribal_motion = tribal_motion;
        self.momentum = momentum;
    }

    /// The number of particles in the swarm. Powers of 2 are recommended
    pub fn num_particles(&mut self, num_particles: usize) {
        self.num_particles = num_particles;
    }
}

impl fmt::Display for SwarmConfig {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} Particles: ", self.num_particles)?;
        write!(f, "Particle motion coefficients: \n")?;
        write!(
            f,
            "\t [Local: {}, tribal: {}] \t momentum: {} \n \t wall_bounce_factor: {} \n",
            self.local_motion, self.tribal_motion, self.momentum, self.wall_bounce_factor,
        )?;
        write!(f, "{} \n", self.global_behavior)?;
        write!(f, "{}", self.transient_behavior)
    }
}

#[derive(Clone, Debug)]
pub enum GlobalBehavior {
    /// collaboration with: (global motion coefficient, collaboration rate)
    Synergic(f64, usize),
    /// no inter-swarm collaboration
    Solitary,
}

impl GlobalBehavior {
    pub fn default() -> Self {
        Self::Synergic(DEFAULT_GLOBAL, DEFAULT_COLLAB_RATE)
    }
}

impl fmt::Display for GlobalBehavior {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Swarm Behavior: ")?;
        match self {
            Self::Synergic(global_motion_coeff, collaboration) => {
                write!(
                    f,
                    "Swarm will collaborate with others every {} iterations \n and weight the global record with a coefficient of {}.",
                    collaboration, global_motion_coeff
                )?;
            }
            Self::Solitary => {
                write!(f, "Swarm will not collaborate with others. \n")?;
            }
        }

        write!(f, "")
    }
}

const DEFAULT_STOCHASTIC_MAG: f32 = 2.0;
const DEFAULT_STOCHASTIC_MODE: u8 = 0;
const DEFAULT_MOMENTUM_MODE: u8 = 0;
const DEFAULT_MOTION_MODE: u8 = 0;

#[derive(Clone, Debug)]
pub struct TransientBehavior {
    pub stochastic_mag: f32,
    pub stochastic_mode: u8,
    pub momentum_mode: u8,
    pub motion_mode: u8,
}

impl TransientBehavior {
    pub fn new() -> Self {
        Self {
            stochastic_mag: DEFAULT_STOCHASTIC_MAG,
            stochastic_mode: DEFAULT_STOCHASTIC_MODE,
            momentum_mode: DEFAULT_MOMENTUM_MODE,
            motion_mode: DEFAULT_MOTION_MODE,
        }
    }

    pub fn stochastic_behavior(&mut self, mag: f32, mode: u8) {
        assert!(mag > 0.0, "Stochastic magnitude must be positive!");
        self.stochastic_mag = mag;
        self.stochastic_mode = mode;
    }

    pub fn momentum_mode(&mut self, mode: u8) {
        self.momentum_mode = mode;
    }

    pub fn motion_mode(&mut self, mode: u8) {
        self.motion_mode = mode;
    }
}

impl fmt::Display for TransientBehavior {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Transient Behavior: \n ")?;
        write!(
            f,
            "Stochastic Mode: {}, \t Momentum Mode: {}, \t Motion Mode: {}",
            self.stochastic_mode, self.momentum_mode, self.motion_mode
        )
    }
}

/// Utility to create a range of similar SwarmConfigs
///
/// Fields in this struct are initialized as the default SwarmConfig values
/// They can either be set to a range with a minimum or maximum value, or a fixed value
pub struct SwarmConfigSampler {
    wall_bounce_factor: FixedRange<f64>,
    local_motion: FixedRange<f64>,
    tribal_motion: FixedRange<f64>,
    global_motion: Option<FixedRange<f64>>,
    global_colab_rate: Option<FixedRange<usize>>,
    momentum: FixedRange<f64>,
    num_particles: FixedRange<usize>,
}

impl SwarmConfigSampler {
    pub fn new() -> Self {
        Self {
            wall_bounce_factor: FixedRange::Fixed(DEFAULT_WALL_BOUNCE),
            local_motion: FixedRange::Fixed(DEFAULT_LOCAL),
            tribal_motion: FixedRange::Fixed(DEFAULT_TRIBAL),
            global_motion: Some(FixedRange::Fixed(DEFAULT_GLOBAL)),
            global_colab_rate: Some(FixedRange::Fixed(DEFAULT_COLLAB_RATE)),
            momentum: FixedRange::Fixed(DEFAULT_MOMENTUM),
            num_particles: FixedRange::Fixed(DEFAULT_NUM_PARTICLES),
        }
    }

    pub fn sample_swarm_config(&self, rng: &mut ThreadRng) -> SwarmConfig {
        let mut sc = SwarmConfig::new();

        sc.wall_bounce(self.wall_bounce_factor.value(rng));

        let local = self.local_motion.value(rng);
        let tribal = self.tribal_motion.value(rng);
        let momentum = self.momentum.value(rng);
        sc.motion_coefficients(local, tribal, momentum);

        if let Some(global_motion_fr) = &self.global_motion {
            let global_colab_rate_fr = self
                .global_colab_rate
                .clone()
                .expect("Should not be None if global_motion is not None");

            let global_motion = global_motion_fr.value(rng);
            let global_colab_rate = global_colab_rate_fr.value(rng);

            sc.synergic_behavior(global_motion, global_colab_rate);
        } else {
            sc.solitary_behavior();
        }

        sc.num_particles(self.num_particles.value(rng));

        sc
    }

    pub fn wall_bounce_range(&mut self, min: f64, max: f64) {
        self.wall_bounce_factor = FixedRange::new_range(min, max)
    }

    pub fn local_motion_range(&mut self, min: f64, max: f64) {
        self.local_motion = FixedRange::new_range(min, max)
    }

    pub fn tribal_motion_range(&mut self, min: f64, max: f64) {
        self.tribal_motion = FixedRange::new_range(min, max)
    }

    pub fn synergic_range(
        &mut self,
        min_global_motion: f64,
        max_global_motion: f64,
        min_colab_rate: usize,
        max_colab_rate: usize,
    ) {
        self.global_motion = Some(FixedRange::new_range(min_global_motion, max_global_motion));
        self.global_colab_rate = Some(FixedRange::new_range(min_colab_rate, max_colab_rate));
    }

    pub fn solitary(&mut self) {
        self.global_motion = None;
        self.global_colab_rate = None;
    }

    pub fn num_particles_range(&mut self, min: usize, max: usize) {
        self.num_particles = FixedRange::new_range(min, max);
    }

    pub fn wall_bounce_fixed(&mut self, value: f64) {
        self.wall_bounce_factor = FixedRange::Fixed(value);
    }

    pub fn local_fixed(&mut self, value: f64) {
        self.local_motion = FixedRange::Fixed(value);
    }

    pub fn tribal_fixed(&mut self, value: f64) {
        self.tribal_motion = FixedRange::Fixed(value);
    }

    pub fn synergic_fized(&mut self, global_value: f64, global_colab_rate: usize) {
        self.global_motion = Some(FixedRange::Fixed(global_value));
        self.global_colab_rate = Some(FixedRange::Fixed(global_colab_rate));
    }
}

#[derive(Clone)]
enum FixedRange<T> {
    Fixed(T),
    Range([T; 2]),
}

impl<T> FixedRange<T>
where
    T: PartialOrd,
{
    pub fn new_range(min: T, max: T) -> Self {
        assert!(min < max, "Minimum value must be less than Maximum value!");
        Self::Range([min, max])
    }
}

trait FRValue<T> {
    fn value(&self, rng: &mut ThreadRng) -> T;
}

impl FRValue<f64> for FixedRange<f64> {
    fn value(&self, rng: &mut ThreadRng) -> f64 {
        match self {
            Self::Fixed(val) => *val,
            Self::Range([min, max]) => rng.gen_range(min, max),
        }
    }
}

impl FRValue<usize> for FixedRange<usize> {
    fn value(&self, rng: &mut ThreadRng) -> usize {
        match self {
            Self::Fixed(val) => *val,
            Self::Range([min, max]) => rng.gen_range(min, max),
        }
    }
}