/*
 * Copyright (c) 2016 Frank Fischer <frank-fischer@shadow-soft.de>
 *
 * This program is free software: you can redistribute it and/or
 * modify it under the terms of the GNU General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see  <http://www.gnu.org/licenses/>
 */

//! Master problem implementation using CPLEX.

use {Real, DVector, Minorant};
use master::UnconstrainedMasterProblem;

use cplex;
use cplex::*;

use std::result;
use std::ffi::CString;
use std::ptr;
use libc::{c_char, c_int};
use std::f64::NEG_INFINITY;
use std::f64;

quick_error! {
    /// A solver error.
    #[derive(Debug)]
    pub enum Error {
        Cplex(err: CplexError) {
            cause(err)
            description(err.description())
            display("{}", err)
            from()
        }

        NoMinorants {
            description("No minorants")
            display("Solver Error: no minorants when solving the master problem")
        }
    }
}

pub type Result<T> = result::Result<T, Error>;

pub struct CplexMaster {
    lp : *mut CPXlp,

    /// True if the QP must be updated.
    force_update: bool,

    /// List of free minorant indices.
    freeinds : Vec<usize>,

    /// List of minorant indices to be updated.
    updateinds: Vec<usize>,

    /// Mapping minorant to index.
    min2index : Vec<Vec<usize>>,

    /// Mapping index to minorant.
    index2min : Vec<(usize, usize)>,

    /// The quadratic term.
    qterm: Vec<DVector>,

    /// The weight of the quadratic term.
    weight: Real,

    /// The minorants for each subproblem in the model.
    minorants: Vec<Vec<Minorant>>,
    /// Optimal multipliers for each subproblem in the model.
    opt_mults: Vec<DVector>,
    /// Optimal aggregated minorant.
    opt_minorant: Minorant,
}

impl Drop for CplexMaster {
    fn drop(&mut self) {
        unsafe { CPXfreeprob(env(), &mut self.lp) };
    }
}


impl UnconstrainedMasterProblem for CplexMaster {
    type Error = Error;

    type MinorantIndex = usize;

    fn new() -> Result<CplexMaster> {
        Ok(CplexMaster {
            lp: ptr::null_mut(),
            force_update: true,
            freeinds: vec![],
            updateinds: vec![],
            min2index: vec![],
            index2min: vec![],
            qterm: vec![],
            weight: 1.0,
            minorants: vec![],
            opt_mults: vec![],
            opt_minorant: Minorant::new(),
        })
    }

    fn num_subproblems(&self) -> usize {
        self.minorants.len()
    }

    fn set_num_subproblems(&mut self, n: usize) -> Result<()> {
        trycpx!(CPXsetintparam(env(), CPX_PARAM_QPMETHOD, CPX_ALG_BARRIER));
        trycpx!(CPXsetdblparam(env(), CPX_PARAM_BAREPCOMP, 1e-12));

        self.min2index = vec![vec![]; n];
        self.index2min.clear();
        self.freeinds.clear();
        self.minorants = vec![vec![]; n];
        self.opt_mults = vec![dvec![]; n];

        Ok(())
    }

    fn weight(&self) -> Real {
        self.weight
    }

    fn set_weight(&mut self, weight: Real) {
        assert!(weight > 0.0);
        self.weight = weight;
    }

    fn num_minorants(&self, fidx : usize) -> usize {
        self.minorants[fidx].len()
    }

    fn add_minorant(&mut self, fidx: usize, minorant: Minorant) -> Result<usize> {
        debug!("Add minorant");
        debug!("  fidx={} index={}: {}", fidx, self.minorants[fidx].len(), minorant);

        let min_idx = self.minorants[fidx].len();
        self.minorants[fidx].push(minorant);
        self.opt_mults[fidx].push(0.0);

        self.force_update = true;

        if let Some(idx) = self.freeinds.pop() {
            self.min2index[fidx].push(idx);
            self.index2min[idx] = (fidx, min_idx);
            self.updateinds.push(idx);
            Ok(idx)
        } else {
            let idx = self.index2min.len();
            self.min2index[fidx].push(idx);
            self.index2min.push((fidx, min_idx));
            self.updateinds.push(idx);
            Ok(idx)
        }
    }

    #[allow(unused_variables)]
    fn solve(&mut self, eta: &DVector, fbound: Real, augbound: Real, relprec: Real) -> Result<()> {
        if self.force_update || !self.updateinds.is_empty() {
            try!(self.init_qp());
        }

        let nvars = unsafe { CPXgetnumcols(env(), self.lp) as usize };
        if nvars == 0 {
            return Err(Error::NoMinorants)
        }
        // update linear costs
        {
            let mut c = Vec::with_capacity(nvars);
            let mut inds = Vec::with_capacity(nvars);
            for mins in self.minorants.iter() {
                for m in mins {
                    inds.push(c.len() as c_int);
                    c.push( -m.constant * self.weight - m.linear.dot(eta));
                }
            }
            trycpx!(CPXchgobj(env(), self.lp, nvars as c_int, inds.as_ptr(), c.as_ptr()));
        }

        trycpx!(CPXqpopt(env(), self.lp));
        let mut sol = vec![0.0; nvars];
        trycpx!(CPXgetx(env(), self.lp, sol.as_mut_ptr(), 0, nvars as c_int -1));

        let mut idx = 0;
        let mut mults = Vec::with_capacity(nvars);
        let mut mins = Vec::with_capacity(nvars);
        for fidx in 0..self.minorants.len() {
            for i in 0..self.minorants[fidx].len() {
                self.opt_mults[fidx][i] = sol[idx];
                mults.push(sol[idx]);
                mins.push(&self.minorants[fidx][i]);
                idx += 1;
            }
        }

        self.opt_minorant.combine_all_ref(&mults, &mins);

        Ok(())
    }

    fn dualopt(&self) -> &DVector {
        &self.opt_minorant.linear
    }

    fn dualopt_cutval(&self) -> Real {
        self.opt_minorant.constant
    }

    fn multiplier(&self, min : usize) -> Real {
        let (fidx, idx) = self.index2min[min];
        self.opt_mults[fidx][idx]
    }

    fn eval_model(&self, y: &DVector) -> Real {
        let mut result = 0.0;
        for mins in &self.minorants {
            let mut this_val = NEG_INFINITY;
            for m in mins {
                this_val = this_val.max(m.eval(y));
            }
            result += this_val;
        }
        result
    }

    fn aggregate(&mut self, fidx: usize, mins: &[usize]) -> Result<(usize, DVector)> {
        assert!(mins.len() > 0, "No minorants specified to be aggregated");

        if mins.len() == 1 { return Ok((mins[0], dvec![1.0])) }

        // scale coefficients
        let mut sum_coeffs = 0.0;
        for &i in mins {
            sum_coeffs += self.opt_mults[fidx][self.index2min[i].1];
        }
        let aggr_coeffs = if sum_coeffs != 0.0 {
            mins.iter().map(|&i| self.opt_mults[fidx][self.index2min[i].1] / sum_coeffs).collect::<DVector>()
        } else {
            dvec![0.0; mins.len()]
        };

        // compute aggregated diagonal term
        let mut aggr_diag = 0.0;
        for (idx_i, &i) in mins.iter().enumerate() {
            for (idx_j, &j) in mins.iter().enumerate() {
                aggr_diag += aggr_coeffs[idx_i] * aggr_coeffs[idx_j] * self.qterm[i][j];
            }
        }

        // compute aggregated off-diagonal terms
        let mut aggr_qterm = dvec![0.0; self.qterm.len()];
        for fidx_i in 0..self.minorants.len() {
            for idx_i in 0..self.minorants[fidx_i].len() {
                let i = self.min2index[fidx_i][idx_i];
                for (idx_j, &j) in mins.iter().enumerate() {
                    aggr_qterm[i] += aggr_coeffs[idx_j] * self.qterm[i][j];
                }
            }
        }

        // aggregate the minorants
        let mut aggr = Minorant::new();
        {
            let mut aggr_mins = Vec::with_capacity(mins.len());

            for &i in mins {
                let (min_fidx, min_idx) = self.index2min[i];
                debug_assert!(min_fidx == fidx);

                let m     = self.minorants[fidx].swap_remove(min_idx);
                let idx   = self.min2index[fidx].swap_remove(min_idx);
                self.opt_mults[fidx].swap_remove(min_idx);
                self.freeinds.push(idx);
                debug_assert!(idx == i);

                aggr_mins.push(m);

                // update index2min table for moved minorant
                if min_idx < self.minorants[fidx].len() {
                    self.index2min[self.min2index[fidx][min_idx]].1 = min_idx;
                }
            }
            aggr.combine_all(&aggr_coeffs, &aggr_mins);
        }

        // save aggregated minorant
        let aggr_idx = self.freeinds.pop().unwrap();
        self.minorants[fidx].push(aggr);
        self.opt_mults[fidx].push(sum_coeffs);
        self.min2index[fidx].push(aggr_idx);
        self.index2min[aggr_idx] = (fidx, self.minorants[fidx].len()-1);

        // update qterm
        for fidx_i in 0..self.minorants.len() {
            for idx_i in 0..self.minorants[fidx_i].len() {
                let i = self.min2index[fidx_i][idx_i];
                self.qterm[i][aggr_idx] = aggr_qterm[i];
                self.qterm[aggr_idx][i] = aggr_qterm[i];
            }
        }
        self.qterm[aggr_idx][aggr_idx] = aggr_diag;

        Ok((aggr_idx, aggr_coeffs))
    }

    fn move_center(&mut self, alpha: Real, d: &DVector) {
        for mins in self.minorants.iter_mut() {
            for m in mins.iter_mut() {
                m.move_center(alpha, d);
            }
        }
    }
}


impl CplexMaster {
    fn init_qp(&mut self) -> Result<()> {
        if self.lp != ptr::null_mut() {
            trycpx!(CPXfreeprob(env(), &mut self.lp));
        }
        trycpx!({
            let mut status = 0;
            self.lp = CPXcreateprob(env(), &mut status, CString::new("mcf").unwrap().as_ptr());
            status
        });

        let nfun = self.minorants.len();
        let mut nvars;

        // add constraints
        {
            let sense: Vec<c_char> = vec!['E' as c_char; nfun];
            let rhs = dvec![1.0; nfun];
            let mut rmatbeg = Vec::with_capacity(nfun);
            let mut rmatind = Vec::with_capacity(self.index2min.len());
            let mut rmatval = Vec::with_capacity(self.index2min.len());

            nvars = 0;
            for i in 0..nfun {
                rmatbeg.push(nvars as c_int);
                for _ in 0..self.minorants[i].len() {
                    rmatind.push(nvars as c_int);
                    rmatval.push(1.0);
                    nvars += 1;
                }
            }

            trycpx!(CPXaddrows(env(), self.lp, nvars as c_int, nfun as c_int, nvars as c_int,
                               rhs.as_ptr(), sense.as_ptr(),
                               rmatbeg.as_ptr(), rmatind.as_ptr(), rmatval.as_ptr(),
                               ptr::null(), ptr::null()));
        }

        // build quadratic term
        {
            self.qterm.resize(self.index2min.len(), dvec![]);
            for i in 0..self.qterm.len() { self.qterm[i].resize(self.index2min.len(), 0.0); }

            // the global indices for each minorant in order
            let mut activeinds = vec![];
            for (fidx, mins_i) in self.minorants.iter().enumerate() {
                for (i, m_i) in mins_i.iter().enumerate() {
                    let idx_i = self.min2index[fidx][i];
                    activeinds.push(idx_i);
                    for &idx_j in &self.updateinds {
                        let (fidx_j, j) = self.index2min[idx_j];
                        let x = m_i.linear.dot(&self.minorants[fidx_j][j].linear);
                        self.qterm[idx_i][idx_j] = x;
                        self.qterm[idx_j][idx_i] = x;
                    }
                }
            }

            for &idx_i in activeinds.iter() {
                for &idx_j in activeinds.iter() {
                    let (fidx_i, i) = self.index2min[idx_i];
                    let (fidx_j, j) = self.index2min[idx_j];
                    let m_i = &self.minorants[fidx_i][i];
                    let m_j = &self.minorants[fidx_j][j];
                    debug_assert!((self.qterm[idx_i][idx_j] - m_i.linear.dot(&m_j.linear)).abs() < 1e-6);
                }
            }

            // main diagonal plus small identity to ensure Q being semi-definite
            let mut maxq = 0.0;
            for &i in &activeinds { maxq = f64::max(maxq, self.qterm[i][i]) }
            maxq *= 1e-8;

            // update coefficients
            for (i, &idx_i) in activeinds.iter().enumerate() {
                for (j, &idx_j) in activeinds.iter().enumerate() {
                    if i != j {
                        trycpx!(CPXchgqpcoef(env(), self.lp, i as c_int, j as c_int, self.qterm[idx_i][idx_j]));
                    } else {
                        trycpx!(CPXchgqpcoef(env(), self.lp, i as c_int, j as c_int, self.qterm[idx_i][idx_j] + maxq));
                    }
                }
            }
        }

        self.updateinds.clear();
        self.force_update = false;

        Ok(())
    }
}
