qopt 0.19.4

//! Defines an abstract optimizer.

use std::{
    cmp::Ordering,
    time::Instant,
};

use maria_linalg::{
    Matrix,
    Vector,
};

pub use super::{
    cli::get_cli,
    error,
    help,
    paradigm::Paradigm,
};

const DX: f64 = 1E-6;

/// Stores information about an optimizer.
pub struct Optimizer<const C: usize, const D: usize, const K: usize> {
    pub print_every: usize,
    paradigm: Paradigm,
    obj: fn(Vector<C>, Vector<D>) -> f64,
    gradient: Option<fn(Vector<C>) -> Vector<C>>,
    hessian: Option<fn(Vector<C>) -> Matrix<C>>,
    criterion: f64,
    maxiter: usize,
    pub maxtemp: f64,
    pub stdev: f64,
}

impl<const C: usize, const D: usize, const K: usize> Optimizer<C, D, K> {
    /// Given a function, construct an `Optimizer`.
    pub fn new(
        obj: fn(Vector<C>, Vector<D>) -> f64,
        gradient: Option<fn(Vector<C>) -> Vector<C>>,
        hessian: Option<fn(Vector<C>) -> Matrix<C>>,
    ) -> Self {
        // Set default values
        let mut print_every = 1;
        let mut paradigm = Paradigm::SteepestDescent;
        let mut criterion = 0.001;
        let mut maxiter = 100;
        let mut maxtemp = 1.0;
        let mut stdev = 1.0;

        get_cli(
            &mut print_every,
            &mut paradigm,
            &mut criterion,
            &mut maxiter,
            &mut maxtemp,
            &mut stdev,
        );

        Self {
            print_every,
            paradigm,
            obj,
            gradient,
            hessian,
            criterion,
            maxiter,
            maxtemp,
            stdev,
        }
    }

    /// Computes the objective of the function to be optimized, accounting for constraints.
    /// 
    /// When a constraint is hit, return `f64::NAN`.
    pub fn objective(&self, continuous: Vector<C>, discrete: Vector<D>) -> f64 {
        (self.obj)(continuous, discrete)
    }

    /// Computes the gradient of the function to be optimized, accounting for constraints.
    /// 
    /// When a constraint is hit, return a vector with each element set to `f64::NAN`.
    pub fn gradient(&self, continuous: Vector<C>, discrete: Vector<D>) -> Vector<C> {
        match self.gradient {
            Some (g) => g(continuous),
            None => {
                let mut gradient = Vector::zero();

                for i in 0..C {
                    let mut high = continuous;
                    let mut low = continuous;

                    high[i] += DX / 2.0;
                    low[i] -= DX / 2.0;

                    let f_high = self.objective(high, discrete);
                    let f_low = self.objective(low, discrete);

                    gradient[i] = (f_high - f_low) / DX;
                }

                gradient
            }
        }
    }

    /// Computes the hessian of the function to be optimized, accounting for constraints.
    /// 
    /// When a constraint is hit, return a matrix with each element set to`f64::NAN`.
    pub fn hessian(&self, continuous: Vector<C>, discrete: Vector<D>) -> Matrix<C> {
        match self.hessian {
            Some (h) => h(continuous),
            None => {
                let mut hessian = Matrix::zero();

                for i in 0..C {
                    let vector = continuous;
                    let mut high = vector;
                    let mut low = vector;

                    high[i] += DX / 2.0;
                    low[i] -= DX / 2.0;

                    let f_high = self.gradient(high, discrete);
                    let f_low = self.gradient(low, discrete);

                    let row = (f_high + f_low.scale(-1.0)).scale(1.0 / DX);

                    for j in 0..C {
                        hessian[(i, j)] = row[j];
                    }
                }

                hessian
            }
        }
    }

    /// Given an input vector, return a step vector.
    pub fn step(&self, continuous: Vector<C>, discrete: Vector<D>) -> Vector<C> {
        self.gradient(continuous, discrete).scale(-1.0)
    }

    /// Given an input population, return an updated (more optimal) population, enforcing discretized values where appropriate.
    pub fn update(&self, iter: usize, continuous: [Vector<C>; K], discrete: [Vector<D>; K], permitted: &[f64]) -> ([Vector<C>; K], [Vector<D>; K]) {
        self.paradigm.update(
            self,
            iter,
            continuous,
            discrete,
            permitted,
        )
    }

    /// Sorts a population vector by increasing objective.
    pub fn sort(&self, continuous: [Vector<C>; K], discrete: [Vector<D>; K]) -> ([Vector<C>; K], [Vector<D>; K]) {
        let mut sorted = [(Vector::zero(), Vector::zero()); K];

        for k in 0..K {
            sorted[k] = (continuous[k], discrete[k]);
        }
        
        sorted.sort_by(|one, two| {
            let a = self.objective(one.0, one.1);
            let b = self.objective(two.0, two.1);

            match (a.is_nan(), b.is_nan()) {
                (true, true) => Ordering::Equal,
                (true, false) => Ordering::Greater,
                (false, true) => Ordering::Less,
                (false, false) => a.partial_cmp(&b).unwrap(),
            }
        });
        
        let mut c = [Vector::zero(); K];
        let mut d = [Vector::zero(); K];

        for k in 0..K {
            (c[k], d[k]) = sorted[k];
        }

        (c, d)
    }

    /// Gets the best element in this population.
    pub fn get_best(&self, continuous: [Vector<C>; K], discrete: [Vector<D>; K]) -> (Vector<C>, Vector<D>) {
        let (c, d) = self.sort(continuous, discrete);
        (c[0], d[0])
    }

    /// Optimize using a mixed approach given an input continuous vector, input discrete vector, and a list of permitted values.
    pub fn optimize(&self, continuous: Vector<C>, discrete: Vector<D>, permitted: &[f64]) -> (Vector<C>, Vector<D>) {
        // Start time
        // Used for computation time
        let start = Instant::now();

        // Iteration count
        // Used to check maxiter condition 
        let mut i = 0;

        // Gradient-based optimization stopping criterion
        let mut criterion = self.gradient(continuous, discrete).norm();

        // Initial population: seeded based on initial value
        let mut c = [Vector::zero(); K];
        let mut d = [Vector::zero(); K];
        for p in 0..K {
            c[p] = match self.paradigm {
                Paradigm::SteepestDescent | Paradigm::Newton => continuous,
                Paradigm::Genetic | Paradigm::SimulatedAnnealing => Vector::<C>::child(&continuous, &continuous, self.stdev),
            };

            d[p] = match self.paradigm {
                Paradigm::SteepestDescent | Paradigm::Newton => discrete,
                Paradigm::Genetic | Paradigm::SimulatedAnnealing => Vector::<D>::child_discrete(&discrete, &discrete, permitted),
            };
        }

        let opt_type = if C == 0 {
            "DISCRETE"
        } else if D == 0 {
            "CONTINUOUS"
        } else {
            "MIXED"
        };

        println!("INITIATING {} OPTIMIZATION", opt_type);
        println!("Continuous variables: {}", C);
        println!("Discrete variables: {}", D);
        println!("Paradigm: {}", self.paradigm);
        println!("Maximum iterations: {}", self.maxiter);
        println!();

        while criterion > self.criterion && i < self.maxiter {
            i += 1;
            let (best_continuous, best_discrete) = self.get_best(c, d);

            if self.print_every != 0 && i % self.print_every == 0 {
                println!("Iteration: {}", i);
                println!("Objective: {:.8}", self.objective(best_continuous, best_discrete));
                println!("Vector");
                if C > 0 {
                    println!("Continuous\n{}", best_continuous);
                }
                if D > 0 {
                    println!("Discrete\n{}", best_discrete);
                }
                println!();
                println!();
            }

            (c, d) = self.update(i, c, d, permitted);
            criterion = self.gradient(best_continuous, best_discrete).norm();
        }

        let time = start.elapsed().as_micros() as f64 / 1000.0;

        if i == self.maxiter {
            println!("Maximum iteration limit reached in {:.3} milliseconds", time);
        } else {
            println!("Convergence achieved in {:.3} milliseconds", time);
        }

        let (best_continuous, best_discrete) = self.get_best(c, d);
        println!("Result\n");
        if C > 0 {
            println!("Continuous\n{}", best_continuous);
        }
        if D > 0 {
            println!("Discrete\n{}", best_discrete);
        }
        println!("Objective: {:.8}", self.objective(best_continuous, best_discrete));
        
        (best_continuous, best_discrete)
    }
}

#[test]
fn sort() {
    let mut sorted = [f64::NAN, 2.0, 1.0, f64::MAX, f64::NAN];

    sorted.sort_by(|a, b| {
        match (a.is_nan(), b.is_nan()) {
            (true, true) => Ordering::Equal,
            (true, false) => Ordering::Greater,
            (false, true) => Ordering::Less,
            (false, false) => a.partial_cmp(&b).unwrap(),
        }
    });

    dbg!(sorted);
}