//! Defines an abstract optimizer.

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

use maria_linalg::Vector;

pub use super::{
    error,
    function::Function,
    paradigm::Paradigm,
};

/// Stores information about an optimizer.
pub struct Optimizer<const N: usize, const K: usize> {
    paradigm: Paradigm,
    pub function: Box<dyn Function<N>>,
    criterion: f64,
    maxiter: usize,
    pub maxtemp: f64,
    pub stdev: f64,
}

impl<const N: usize, const K: usize> Optimizer<N, K> {
    /// Get command-line arguments.
    fn get_cli(
        paradigm: &mut Paradigm,
        criterion: &mut f64,
        maxiter: &mut usize,
        maxtemp: &mut f64,
        stdev: &mut f64,
    ) {
        // Read in command-line arguments
        let args = env::args().collect::<Vec<String>>();
        let mut i = 1;

        while i < args.len() {
            let arg = args[i].as_str();

            match arg {
                "--paradigm" => {
                    i += 1;
                    if i == args.len() {
                        error("Missing paradigm");
                    }

                    // Set paradigm
                    *paradigm = match args[i].as_str() {
                        "steepest-descent" => Paradigm::SteepestDescent,
                        "newton" => Paradigm::Newton,
                        "genetic" => Paradigm::Genetic,
                        "simulated-annealing" => Paradigm::SimulatedAnnealing,
                        _ => error("Unrecognized paradigm"),
                    };
                    
                    i += 1;
                },
                "--criterion" => {
                    i += 1;
                    if i == args.len() {
                        error("Missing criterion");
                    }

                    // Set criterion
                    *criterion = match str::parse::<f64>(&args[i]) {
                        Ok (m) => m,
                        Err (_) => error("Could not parse as floating-point value"),
                    };
                    
                    i += 1;
                },
                "--maxiter" => {
                    i += 1;
                    if i == args.len() {
                        error("Missing maxiter");
                    }

                    // Set maxiter
                    *maxiter = match str::parse::<usize>(&args[i]) {
                        Ok (m) => m,
                        Err (_) => error("Could not parse as integer value"),
                    };
                    
                    i += 1;
                },
                "--maxtemp" => {
                    i += 1;
                    if i == args.len() {
                        error("Missing maximum annealing temperature");
                    }

                    // Set maxtemp
                    *maxtemp = match str::parse::<f64>(&args[i]) {
                        Ok (m) => m,
                        Err (_) => error("Could not parse as floating-point value"),
                    };
                    
                    i += 1;
                },
                "--stdev" => {
                    i += 1;
                    if i == args.len() {
                        error("Missing standard deviation");
                    }

                    // Set stdev
                    *stdev = match str::parse::<f64>(&args[i]) {
                        Ok (m) => m,
                        Err (_) => error("Could not parse as floating-point value"),
                    };
                    
                    i += 1;
                },
                _ => error("Unrecognized command line argument"),
            }
        }
    }

    /// Given a paradigm, function, and stopping criterion, construct an Optimizer.
    pub fn new(
        function: Box<dyn Function<N>>,
    ) -> Self {
        // Set default values
        let mut paradigm = Paradigm::SteepestDescent;
        let mut criterion = 0.001;
        let mut maxiter = 100;
        let mut maxtemp = 1.0;
        let mut stdev = 1.0;

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

        Self {
            paradigm,
            function,
            criterion,
            maxiter,
            maxtemp,
            stdev,
        }
    }

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

    /// Given an input population, return an updated (more optimal) population.
    pub fn update(&self, iter: usize, population: [Vector<N>; K]) -> [Vector<N>; K] {
        self.paradigm.update(self, iter, population)
    }

    /// Sorts a population vector by decreasing objective.
    pub fn sort(&self, population: [Vector<N>; K]) -> [Vector<N>; K] {
        let mut sorted = population;
        
        sorted.sort_by(
            |a, b| self.function.objective(*a)
                .partial_cmp(
                    &self.function.objective(*b)
                ).unwrap_or(Ordering::Less)
        );
        sorted.reverse();
        
        sorted
    }

    /// Gets the best element in this population.
    pub fn get_best(&self, population: [Vector<N>; K]) -> Vector<N> {
        self.sort(population)[0]
    }

    /// Run the optimizer on a given input vector.
    pub fn optimize(&self, input: Vector<N>) -> Vector<N> {
        // 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.function.gradient(input).norm();

        // Initial population
        let mut population = [input; K];

        println!("INITIATING OPTIMIZATION");
        println!("Paradigm: {}", self.paradigm);
        println!("Stopping criterion: {}", self.criterion);
        println!("Maximum iterations: {}", self.maxiter);
        println!();

        while criterion > self.criterion && i < self.maxiter {
            i += 1;
            let best = self.get_best(population);
            println!("Iteration: {}", i);
            println!("Objective: {:.8}", self.function.objective(best));
            println!("Vector\n{}", best);
            println!("Gradient magnitude: {:.8}", criterion);
            population = self.update(i, population);
            criterion = self.function.gradient(best).norm();
            println!();
            println!();
        }

        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 = self.get_best(population);
        println!("Result\n{}", best);
        println!("Objective: {:.8}", self.function.objective(best));
        
        best
    }
}