// Copyright (c) 2016, 2017, 2018, 2019 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/>
//

pub mod unconstrained;
use self::unconstrained::UnconstrainedMasterProblem;

use super::MasterProblem;
pub use super::SubgradientExtension;
use crate::{DVector, Minorant, Real};

use either::Either;
use itertools::multizip;
use log::debug;
use std::f64::{EPSILON, INFINITY, NEG_INFINITY};

/**
 * Turn unconstrained master problem into box-constrained one.
 *
 * This master problem adds box constraints to an unconstrainted
 * master problem implementation. The box constraints are enforced by
 * an additional outer optimization loop.
 */
pub struct BoxedMasterProblem<M> {
    /// The unconstrained master problem solver.
    pub master: M,

    /// Maximal number of updates of box multipliers.
    pub max_updates: usize,

    /// Variable lower bounds.
    lb: DVector,

    /// Variable upper bounds.
    ub: DVector,

    /// Multipliers for box constraints.
    eta: DVector,

    /// Primal optimal solution.
    primopt: DVector,

    /// Primal optimal solution value.
    primoptval: Real,

    /// Square of norm of dual optimal solution.
    dualoptnorm2: Real,

    /// Model precision.
    model_eps: Real,

    need_new_candidate: bool,

    /// Current number of updates.
    cnt_updates: usize,
}

impl<M> BoxedMasterProblem<M>
where
    M: UnconstrainedMasterProblem,
{
    pub fn with_master(master: M) -> BoxedMasterProblem<M> {
        BoxedMasterProblem {
            master,
            max_updates: 50,
            lb: dvec![],
            ub: dvec![],
            eta: dvec![],
            primopt: dvec![],
            primoptval: 0.0,
            dualoptnorm2: 0.0,
            model_eps: 0.6,
            cnt_updates: 0,
            need_new_candidate: true,
        }
    }

    pub fn set_max_updates(&mut self, max_updates: usize) -> Result<(), M::Err> {
        assert!(max_updates > 0);
        self.max_updates = max_updates;
        Ok(())
    }

    /**
     * Update box multipliers $\eta$.
     *
     * This function solves the dual problem for fixed aggregated
     * minorant w.r.t. $\eta$. When called, the variable `self.primopt`
     * should contain the primal solution (i.e. the new candidate point
     * d) without the influence of $\eta$.
     */
    fn update_box_multipliers(&mut self) -> bool {
        let mut updated_eta = false;
        let weight = self.master.weight();

        self.eta.resize(self.lb.len(), 0.0);
        for (&lb, &ub, x, eta) in multizip((
            self.lb.iter(),
            self.ub.iter(),
            self.primopt.iter_mut(),
            self.eta.iter_mut(),
        )) {
            let newx = if *x < lb {
                lb
            } else if *x > ub {
                ub
            } else {
                *eta = 0.0;
                continue;
            };
            let neweta = (newx - *x) * weight;
            updated_eta = updated_eta || (neweta - *eta).abs() > EPSILON;
            *eta = neweta;
            *x = newx;
        }

        debug!("Eta update");
        debug!("  primopt={}", self.primopt);
        debug!("  eta    ={}", self.eta);

        updated_eta
    }

    // Compute the new candidate point.
    //
    // This consists of two steps:
    //
    // 1. the new point is computed as $-\tfrac{1}{u}\bar{g}$, where $\bar{g}$
    //    is the aggregated minorant
    // 2. the multipliers $\eta$ are updated
    //
    // In other words, this function computes the new candidate
    // defined by a fixed $\bar{g}$ while choosing the best possible
    // $\eta$.
    //
    fn compute_candidate(&mut self) {
        self.need_new_candidate = false;

        if self.master.dualopt().len() == self.lb.len() {
            self.primopt.scal(-1.0 / self.master.weight(), self.master.dualopt())
        } else {
            self.primopt.init0(self.lb.len());
        }
        self.update_box_multipliers();
    }

    /// Compute $\langle b, \eta \rangle$ with $b$ the bounds of eta.
    fn eta_cutval(&self) -> Real {
        multizip((self.lb.iter(), self.ub.iter(), self.eta.iter()))
            .map(|(&lb, &ub, &eta)| {
                if eta >= 0.0 {
                    if lb > NEG_INFINITY {
                        lb * eta
                    } else {
                        0.0
                    }
                } else if ub < INFINITY {
                    ub * eta
                } else {
                    0.0
                }
            })
            .sum()
    }

    /**
     * Return $\\|G \alpha - \eta\\|_2\^2$.
     *
     * This is the norm-square of the dual optimal solution including
     * the current box-multipliers $\eta$.
     */
    fn get_norm_subg2(&self) -> Real {
        self.eta.dot(self.master.dualopt())
    }
}

impl<M> MasterProblem for BoxedMasterProblem<M>
where
    M: UnconstrainedMasterProblem,
{
    type MinorantIndex = M::MinorantIndex;

    type Err = M::Err;

    fn new() -> Result<Self, Self::Err> {
        M::new().map(BoxedMasterProblem::with_master)
    }

    fn set_num_subproblems(&mut self, n: usize) -> Result<(), Self::Err> {
        self.master.set_num_subproblems(n)
    }

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

    fn num_variables(&self) -> usize {
        self.lb.len()
    }

    fn compress<F>(&mut self, f: F) -> Result<(), Self::Err>
    where
        F: FnMut(Self::MinorantIndex, &mut dyn Iterator<Item = (Self::MinorantIndex, Real)>),
    {
        self.master.compress(f)
    }

    fn fill_models<F>(&mut self, f: F) -> Result<(), Self::Err>
    where
        F: FnMut(Self::MinorantIndex, &mut dyn Iterator<Item = (Self::MinorantIndex, Real)>),
    {
        self.master.fill_models(f)
    }

    fn set_vars(&mut self, n: usize, lb: Option<DVector>, ub: Option<DVector>) -> Result<(), Self::Err> {
        assert_eq!(lb.as_ref().map(|x| x.len()).unwrap_or(n), n);
        assert_eq!(ub.as_ref().map(|x| x.len()).unwrap_or(n), n);
        self.lb = lb.unwrap_or_else(|| dvec![NEG_INFINITY; n]);
        self.ub = ub.unwrap_or_else(|| dvec![INFINITY; n]);
        Ok(())
    }

    fn num_minorants(&self, fidx: usize) -> usize {
        self.master.num_minorants(fidx)
    }

    fn add_minorant(&mut self, fidx: usize, minorant: Minorant) -> Result<Self::MinorantIndex, Self::Err> {
        self.master.add_minorant(fidx, minorant)
    }

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

    fn set_weight(&mut self, weight: Real) -> Result<(), Self::Err> {
        self.master.set_weight(weight)
    }

    fn add_vars<S>(
        &mut self,
        bounds: &[(Option<usize>, Real, Real)],
        extend_subgradient: S,
    ) -> Result<(), Either<Self::Err, S::Err>>
    where
        S: SubgradientExtension<Self::MinorantIndex>,
    {
        if !bounds.is_empty() {
            for (index, l, u) in bounds.iter().filter_map(|v| v.0.map(|i| (i, v.1, v.2))) {
                self.lb[index] = l;
                self.ub[index] = u;
            }
            self.lb.extend(bounds.iter().filter(|v| v.0.is_none()).map(|x| x.1));
            self.ub.extend(bounds.iter().filter(|v| v.0.is_none()).map(|x| x.2));
            self.eta.resize(self.lb.len(), 0.0);
            self.need_new_candidate = true;
            let nnew = bounds.iter().filter(|v| v.0.is_none()).count();
            let changed = bounds.iter().filter_map(|v| v.0).collect::<Vec<_>>();
            self.master.add_vars(nnew, &changed, extend_subgradient)
        } else {
            Ok(())
        }
    }

    #[allow(clippy::cognitive_complexity)]
    fn solve(&mut self, center_value: Real) -> Result<(), Self::Err> {
        debug!("Solve Master");
        debug!("  lb      ={}", self.lb);
        debug!("  ub      ={}", self.ub);

        if self.need_new_candidate {
            self.compute_candidate();
        }

        let mut cnt_updates = 0;
        let mut old_augval = NEG_INFINITY;
        loop {
            cnt_updates += 1;
            self.cnt_updates += 1;

            // TODO: relprec is fixed
            self.master.solve(&self.eta, center_value, old_augval, 1e-3)?;

            // compute the primal solution without the influence of eta
            self.primopt.scal(-1.0 / self.master.weight(), self.master.dualopt());

            // solve w.r.t. eta
            let updated_eta = self.update_box_multipliers();

            // compute value of the linearized model
            self.dualoptnorm2 = self.get_norm_subg2();
            let linval = self.master.dualopt().dot(&self.primopt) + self.master.dualopt_cutval();
            assert!((self.eta.dot(&self.primopt) - self.eta_cutval()).abs() < 1e-6); // verify complementarity
            let augval = linval + 0.5 * self.dualoptnorm2 / self.master.weight();

            let mut cutval = linval;
            if updated_eta {
                cutval = self.master.eval_model(&self.primopt)
            }

            let curval = center_value;
            let model_prec = (cutval - linval) / (curval - linval).max(1e-16);

            debug!("Eta Test");
            debug!("  dualnorm2={}", self.dualoptnorm2);
            debug!("  linval={}", linval);
            debug!("  modval={}", self.master.eval_model(&self.primopt));
            debug!("  augval={}", augval);
            debug!("  cutval={}", cutval);
            debug!("  model_prec={}", model_prec);
            debug!("  old_augval={}", old_augval);
            debug!("  center_value={}", center_value);
            debug!("  model_eps={}", self.model_eps);
            debug!(
                "  cut-lin={} < eps*(cur-lin)={}",
                cutval - linval,
                self.model_eps * (curval - linval)
            );
            debug!("  cnt_update={} max_updates={}", cnt_updates, self.max_updates);

            self.primoptval = linval;

            if augval < old_augval + 1e-10
                || cutval - linval < self.model_eps * (curval - linval)
                || cnt_updates >= self.max_updates
            {
                break;
            }

            old_augval = old_augval.max(augval);
        }

        debug!("Model");
        debug!("  cnt_update={}", cnt_updates);
        debug!("  primopt={}", self.primopt);
        debug!("  dualopt={}", self.master.dualopt());
        debug!("  etaopt={}", self.eta);
        debug!("  primoptval={}", self.primoptval);

        Ok(())
    }

    fn aggregate(&mut self, fidx: usize, mins: &[usize]) -> Result<(Self::MinorantIndex, DVector), Self::Err> {
        self.master.aggregate(fidx, mins)
    }

    fn get_primopt(&self) -> DVector {
        self.primopt.clone()
    }

    fn get_primoptval(&self) -> Real {
        self.primoptval
    }

    fn get_dualoptnorm2(&self) -> Real {
        self.dualoptnorm2
    }

    fn multiplier(&self, min: Self::MinorantIndex) -> Real {
        self.master.multiplier(min)
    }

    fn opt_multipliers<'a>(&'a self, fidx: usize) -> Box<dyn Iterator<Item = (Self::MinorantIndex, Real)> + 'a> {
        self.master.opt_multipliers(fidx)
    }

    fn move_center(&mut self, alpha: Real, d: &DVector) {
        self.need_new_candidate = true;
        self.master.move_center(alpha, d);
        self.lb.add_scaled_begin(-alpha, d);
        self.ub.add_scaled_begin(-alpha, d);
    }

    fn set_max_updates(&mut self, max_updates: usize) -> Result<(), Self::Err> {
        BoxedMasterProblem::set_max_updates(self, max_updates)
    }

    fn cnt_updates(&self) -> usize {
        self.cnt_updates
    }
}

/// Builder for `BoxedMasterProblem`.
///
/// `B` is a builder of the underlying `UnconstrainedMasterProblem`.
#[derive(Default)]
pub struct Builder<B>(B);

impl<B> super::Builder for Builder<B>
where
    B: unconstrained::Builder,
    B::MasterProblem: UnconstrainedMasterProblem,
{
    type MasterProblem = BoxedMasterProblem<B::MasterProblem>;

    fn build(&mut self) -> Result<Self::MasterProblem, <Self::MasterProblem as MasterProblem>::Err> {
        self.0.build().map(BoxedMasterProblem::with_master)
    }
}

impl<B> std::ops::Deref for Builder<B> {
    type Target = B;

    fn deref(&self) -> &B {
        &self.0
    }
}

impl<B> std::ops::DerefMut for Builder<B> {
    fn deref_mut(&mut self) -> &mut B {
        &mut self.0
    }
}
