lbfgsbrs 0.1.1

use super::algorithm::setulb;
use super::timer::timer;

use log::info;

pub const START: i32 = 1;
pub const NEW_X: i32 = 2;
pub const ABNORMAL: i32 = 3; 
pub const FG: i32 = 10;
pub const FG_END: i32 = 15;
pub const CONV_GRAD: i32 = 21; 
pub const CONV_F: i32 = 22; 
pub const STOP: i32 = 30;
pub const STOP_END: i32 = 40;
pub const STOP_ITER: i32 = 32;
pub const STOP_GRAD: i32 = 33;
pub const ERROR: i32 = 200; 
pub const ERROR_END: i32 = 240; 

/// Parameters for the L-BFGS-B optimization algorithm
pub struct LbfgsbParameters {
    /// The maximum number of stored L-BFGS iteration pairs
    pub m: usize,
    /// The convergence tolerance for the projected gradient
    pub pgtol: f64,
    /// The convergence tolerance factor for function value
    pub factr: f64,
    /// Time limit in seconds (0.0 means no limit)
    pub time_limit: f64,
    /// Maximum number of iterations (0 means use default)
    pub max_iter: i32,
}

impl Default for LbfgsbParameters {
    fn default() -> Self {
        LbfgsbParameters {
            m: 10,
            pgtol: 1e-5,
            factr: 1e7,
            time_limit: 0.2,
            max_iter: 1000,
        }
    }
}

/// L-BFGS-B minimizer for bound-constrained optimization
pub struct LbfgsbMinimizer<'a, F, G>
where
    F: Fn(&Vec<f64>) -> f64,
    G: Fn(&Vec<f64>) -> Vec<f64>,
{
    n: usize,

    // Current solution vector
    x: &'a mut Vec<f64>,

    // Objective function
    f: &'a F,

    // Gradient function
    g: &'a G,

    // Lower bounds
    l: Vec<f64>,

    // Upper bounds
    u: Vec<f64>,
    
    // Compatibility 
    nbd: Vec<i32>,
    wa: Vec<f64>,
    iwa: Vec<i32>,
    task: i32,
    csave: i32,
    lsave: Vec<bool>,
    isave: Vec<i32>,
    dsave: Vec<f64>,

    params: LbfgsbParameters
}

impl<'a, F, G> LbfgsbMinimizer<'a, F, G>
where
    F: Fn(&Vec<f64>) -> f64,
    G: Fn(&Vec<f64>) -> Vec<f64>,
{
    /// Create a new L-BFGS-B minimizer
    pub fn new(
        x0: &'a mut Vec<f64>,
        f: &'a F,
        g: &'a G,
        params: Option<LbfgsbParameters>,
    ) -> Self {
        let n = x0.len();
        let params = params.unwrap_or_default();
        
        // Default bounds
        let l = vec![f64::NEG_INFINITY; n];
        let u = vec![f64::INFINITY; n];
        
        // Creating lbfgs struct with parameters or defaults
        LbfgsbMinimizer {
            n: n,
            x: x0,
            l: vec![0.0f64;n as usize],
            u: vec![0.0f64;n as usize],
            nbd: vec![0;n as usize],
            f: f,
            g: g,
            wa: vec![0.0f64;(2*params.m*n+11*params.m*params.m+params.m*n+8*params.m) as usize],
            iwa: vec![0;3*n as usize],
            task: 0,
            csave: 0,
            lsave: vec![false;4],
            isave: vec![0;44],
            dsave: vec![0.0f64;30],
            params: params
        }
    }

    // this function will start the optimization algorithm
    pub fn minimize(&mut self) {
        // Initialize objective variables
        let mut f_val = (self.f)(self.x);
        let mut g_val = (self.g)(self.x);

        self.task = START;

        let time_begin = timer();

        // start of the loop
        loop {
            setulb(self.n,
                self.params.m,
                self.x.as_mut_slice(),
                self.l.as_mut_slice(),
                self.u.as_mut_slice(),
                self.nbd.as_mut_slice(),
                &mut f_val,
                g_val.as_mut_slice(),
                self.params.factr,
                self.params.pgtol,
                self.wa.as_mut_slice(),
                self.iwa.as_mut_slice(),
                &mut self.task,
                &mut self.csave,
                self.lsave.as_mut_slice(),
                self.isave.as_mut_slice(),
                self.dsave.as_mut_slice()
            );
            
            if FG <= self.task && self.task <= FG_END {
                f_val = (self.f)(self.x);
                g_val = (self.g)(self.x);
            }
            else if self.task == NEW_X {
                // Check stopping criteria
                if self.isave[33] >= self.params.max_iter {
                    self.task = STOP_ITER;
                }
                
                if self.dsave[12] <= (f_val.abs() + 1.0) * 1e-10 {
                    self.task = STOP_GRAD;
                }
                
                info!("Iterate {}  nfg = {}   f = {}   |proj g| = {}", 
                    self.isave[29], self.isave[33], f_val, self.dsave[12]);
            }
            else {
                // Exit the loop if task is not FG, FG_END, or NEW_X
                break;
            }
        }
    }
    // this function returns the solution after minimization
    pub fn get_x(&self) -> Vec<f64> {
        self.x.clone()
    }
    // this function is used to set lower bounds to a variable
    pub fn set_lower_bound(&mut self, index: usize, value: f64) {
        if self.nbd[index] == 1 || self.nbd[index] == 2 {
            println!("Variable already has Lower Bound");
        } else {
            let temp = self.nbd[index] - 1;
            self.nbd[index] = if temp < 0 {
                -temp
            } else {
                temp
            };
            self.l[index] = value;
        }
    }
    // this function is used to set upper bounds to a variable
    pub fn set_upper_bound(&mut self, index: usize, value: f64) {
        if self.nbd[index] == 3 || self.nbd[index] == 2 {
            println!("Variable already has Lower Bound");
        } else {
            self.nbd[index] = 3 - self.nbd[index];
            self.u[index] = value;
        }
    }

    // set termination tolerance
    // 1.0e12 for low accuracy
    // 1.0e7  for moderate accuracy
    // 1.0e1  for extremely high accuracy
    pub fn set_termination_tolerance(&mut self, t: f64) {
        self.params.factr = t;
    }
    // set tolerance of projection gradient
    pub fn set_tolerance(&mut self, t: f64) {
        self.params.pgtol = t;
    }
    // set max iteration
    pub fn max_iteration(&mut self, i: i32) {
        self.params.max_iter = i;
    }
    // set maximum number of variable metric corrections
    // The range  3 <= m <= 20 is recommended
    pub fn set_metric_correction(&mut self, m: i32) {
        self.params.m = m as usize;
    }
}