// Copyright (c) 2016-2022 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/>
//

use super::UnconstrainedMasterProblem;
use crate::problem::SubgradientExtender;
use crate::saveable::Saveable;
use crate::{Aggregatable, DVector, Minorant, Real};

use either::Either;
use log::debug;
use thiserror::Error;

use std::convert::Infallible;
use std::error::Error;
use std::f64::NEG_INFINITY;

#[cfg(feature = "serialize")]
use serde::{Deserialize, Serialize};

/// Minimal master problem error.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum MinimalMasterError {
    #[error("master problem contains no minorants")]
    NoMinorants,
}

impl AsRef<dyn Error + 'static> for MinimalMasterError {
    fn as_ref(&self) -> &(dyn Error + 'static) {
        self
    }
}

/**
 * A minimal master problem with only two minorants.
 *
 * This is the simplest possible master problem for bundle methods. It
 * has only two minorants and only one function model. The advantage
 * is that this model can be solved explicitely and very quickly, but
 * it is only a very loose approximation of the objective function.
 *
 * Because of its properties, it can only be used if the problem to be
 * solved has a maximal number of minorants of two and only one
 * subproblem.
 */
#[cfg_attr(feature = "serialize", derive(Serialize, Deserialize))]
#[derive(Debug, Clone)]
pub struct MinimalMaster<Mn: Minorant> {
    /// The weight of the quadratic term.
    weight: Real,

    /// The minorants in the model.
    ///
    /// There are up to two minorants, each minorant consists of one part for
    /// each subproblem.
    ///
    /// The minorants for the i-th subproblem have the indices `2*i` and
    /// `2*i+1`.
    minorants: [Vec<Mn>; 2],

    /// Additional minorants for each subproblem.
    ///
    /// These minorants will not be used but removed at the next compression.
    /// Because they will never be used, we throw them away and store only
    /// their index.
    rest_minorants: Vec<Vec<usize>>,

    /// The is the total number of rest minorants.
    ///
    /// Just to keep the minorant indices unique.
    num_rest_minorants: usize,

    /// The number of minorants. Only the minorants with index less than this
    /// number are valid.
    num_minorants: usize,
    /// The number of minorants for each subproblem.
    num_minorants_of: Vec<usize>,
    /// The number of subproblems.
    num_subproblems: usize,
    /// The number of subproblems with at least 1 minorant.
    num_subproblems_with_1: usize,
    /// The number of subproblems with at least 2 minorants.
    num_subproblems_with_2: usize,
    /// Optimal multipliers.
    opt_mult: [Real; 2],
    /// Optimal aggregated minorant.
    opt_minorant: (Real, DVector),
}

impl<Mn: Minorant> UnconstrainedMasterProblem<Mn> for MinimalMaster<Mn> {
    type MinorantIndex = usize;

    type Err = MinimalMasterError;

    fn new() -> Result<MinimalMaster<Mn>, Self::Err> {
        Ok(MinimalMaster {
            weight: 1.0,
            num_minorants: 0,
            num_minorants_of: vec![],
            num_subproblems: 0,
            num_subproblems_with_1: 0,
            num_subproblems_with_2: 0,
            minorants: [vec![], vec![]],
            rest_minorants: vec![],
            num_rest_minorants: 0,
            opt_mult: [0.0, 0.0],
            opt_minorant: Default::default(),
        })
    }

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

    fn set_num_subproblems(&mut self, n: usize) -> Result<(), Self::Err> {
        self.num_subproblems = n;
        self.num_minorants = 0;
        self.num_minorants_of = vec![0; n];
        self.num_subproblems_with_1 = 0;
        self.num_subproblems_with_2 = 0;
        self.minorants = [
            (0..n).map(|_| Default::default()).collect(),
            (0..n).map(|_| Default::default()).collect(),
        ];
        self.rest_minorants = vec![vec![]; n];
        self.num_rest_minorants = 0;
        Ok(())
    }

    fn compress(&mut self) -> Result<(), Self::Err> {
        if self.num_minorants == 2 {
            debug!("Aggregate");
            debug!(
                "  {} * {:?}",
                self.opt_mult[0],
                self.minorants[0].iter().map(|m| m.debug()).collect::<Vec<_>>()
            );
            debug!(
                "  {} * {:?}",
                self.opt_mult[1],
                self.minorants[1].iter().map(|m| m.debug()).collect::<Vec<_>>()
            );

            for fidx in 0..self.num_subproblems {
                self.rest_minorants[fidx].clear();
            }

            self.minorants[0] = Aggregatable::combine(self.opt_mult.iter().cloned().zip(&self.minorants));
            self.opt_mult[0] = 1.0;
            self.num_minorants = 1;
            self.num_minorants_of.clear();
            self.num_minorants_of.resize(self.num_subproblems, 1);
            self.num_subproblems_with_2 = 0;
            self.num_rest_minorants = 0;

            debug!(
                "  {:?}",
                self.minorants[0].iter().map(|m| m.debug()).collect::<Vec<_>>()
            );
        }
        Ok(())
    }

    fn fill_models<F>(&mut self, mut f: F) -> Result<(), Self::Err>
    where
        F: FnMut(Self::MinorantIndex, &mut dyn Iterator<Item = (Self::MinorantIndex, Real)>),
    {
        if self.num_minorants < 2 && self.num_subproblems_with_2 > 0 {
            // there are problems with too few minorants
            for fidx in 0..self.num_subproblems {
                if self.num_minorants_of[fidx] == 1 {
                    let newindex = self.add_minorant(fidx, self.minorants[0][fidx].clone())?;
                    f(newindex, &mut std::iter::once((2 * fidx, 1.0)));
                }
            }
        }
        Ok(())
    }

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

    fn set_weight(&mut self, weight: Real) -> Result<(), Self::Err> {
        assert!(weight > 0.0);
        self.weight = weight;
        Ok(())
    }

    fn num_minorants(&self, fidx: usize) -> usize {
        // TODO: is this correct? What about the pending minorants?
        self.num_minorants_of[fidx]
    }

    fn add_minorant(&mut self, fidx: usize, minorant: Mn) -> Result<usize, Self::Err> {
        // if we have too many minorants, throw them away
        if self.num_minorants_of[fidx] >= 2 {
            debug!("Additional minorant for subproblem {} ignored", fidx);
            let index = 2 * self.num_subproblems + self.num_rest_minorants;
            self.num_rest_minorants += 1;
            self.rest_minorants[fidx].push(index);
            return Ok(index);
        }

        let minidx = self.num_minorants_of[fidx];
        self.num_minorants_of[fidx] += 1;
        self.minorants[minidx][fidx] = minorant;

        match minidx {
            0 => {
                self.num_subproblems_with_1 += 1;
                if self.num_subproblems_with_1 == self.num_subproblems {
                    self.num_minorants = 1;
                    self.opt_mult[0] = 0.0;
                }
                Ok(2 * fidx)
            }
            1 => {
                self.num_subproblems_with_2 += 1;
                if self.num_subproblems_with_2 == self.num_subproblems {
                    self.num_minorants = 2;
                    self.opt_mult[1] = 0.0;
                }
                Ok(2 * fidx + 1)
            }
            _ => unreachable!("Invalid number of minorants in subproblem {}", fidx),
        }
    }

    fn add_vars<SErr>(
        &mut self,
        nnew: usize,
        sgext: &dyn SubgradientExtender<Mn, SErr>,
    ) -> Result<(), Either<Self::Err, SErr>> {
        if nnew == 0 || self.num_subproblems_with_1 == 0 {
            return Ok(());
        }

        for fidx in 0..self.num_subproblems {
            for i in 0..self.num_minorants_of[fidx] {
                sgext
                    .extend_subgradient(fidx, &mut self.minorants[i][fidx])
                    .map_err(Either::Right)?;
            }
        }

        Ok(())
    }

    #[allow(unused_variables)]
    fn solve(&mut self, eta: &DVector, fbound: Real, augbound: Real, relprec: Real) -> Result<(), Self::Err> {
        for fidx in 0..self.num_subproblems {
            for i in 0..self.num_minorants_of[fidx] {
                debug!("  min(fidx:{}, i:{}) = {:?}", fidx, i, self.minorants[i][fidx].debug());
            }
        }

        if self.num_minorants == 2 {
            let mut min0 = (0.0, dvec![]);
            let mut min1 = (0.0, dvec![]);

            Mn::combine_to_vec(
                (0..self.num_subproblems).map(|fidx| (1.0, &self.minorants[0][fidx])),
                &mut min0,
            );
            Mn::combine_to_vec(
                (0..self.num_subproblems).map(|fidx| (1.0, &self.minorants[1][fidx])),
                &mut min1,
            );

            let xx = min0.product(&min0);
            let yy = min1.product(&min1);
            let xy = min0.product(&min1);
            let xeta = min0.linear(eta);
            let yeta = min1.linear(eta);
            let a = yy - 2.0 * xy + xx;
            let b = xy - xx - yeta + xeta;

            let mut alpha2 = 0.0;
            if a > 0.0 {
                alpha2 = ((min1.constant() - min0.constant()) * self.weight - b) / a;
                alpha2 = alpha2.max(0.0).min(1.0);
            }
            self.opt_mult[0] = 1.0 - alpha2;
            self.opt_mult[1] = alpha2;
            <(f64, DVector)>::combine_to_vec(
                self.opt_mult.iter().cloned().zip([min0, min1].iter()),
                &mut self.opt_minorant,
            );
        } else if self.num_minorants == 1 {
            let mins0 = &self.minorants[0];
            Mn::combine_to_vec(
                std::iter::repeat(1.0).zip(mins0).take(self.num_subproblems),
                &mut self.opt_minorant,
            );
            self.opt_mult[0] = 1.0;
        } else {
            return Err(MinimalMasterError::NoMinorants);
        }

        debug!("Unrestricted");
        debug!("  opt_minorant={:?}", self.opt_minorant.debug());
        debug!("  opt_mult={:?}", &self.opt_mult[0..self.num_minorants]);

        Ok(())
    }

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

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

    fn multiplier(&self, min: usize) -> Real {
        self.opt_mult[min % 2]
    }

    fn aggregated_primal(&self, fidx: usize) -> Result<Mn::Primal, Self::Err> {
        Ok(Aggregatable::combine(
            self.opt_mult
                .iter()
                .take(self.num_minorants_of[fidx])
                .enumerate()
                .map(|(i, alpha)| (*alpha, self.minorants[i][fidx].primal())),
        ))
    }

    fn eval_model(&self, y: &DVector) -> Real {
        let mut result = NEG_INFINITY;
        for mins in &self.minorants[0..self.num_minorants] {
            result = result.max(mins.iter().map(|m| m.eval(y)).sum());
        }

        debug_assert!(
            result
                <= (0..self.num_subproblems())
                    .map(|i| self.eval_model_subproblem(i, y))
                    .sum::<Real>()
        );

        result
    }

    fn eval_model_subproblem(&self, fidx: usize, y: &DVector) -> Real {
        let mut result = NEG_INFINITY;
        for mins in &self.minorants[0..self.num_minorants] {
            result = result.max(mins[fidx].eval(y));
        }
        result
    }

    fn move_center(&mut self, alpha: Real, d: &DVector) {
        for fidx in 0..self.num_subproblems {
            for i in 0..self.num_minorants_of[fidx] {
                self.minorants[i][fidx].move_center_begin(alpha, d);
            }
        }
    }
}

pub struct Builder;

impl Default for Builder {
    fn default() -> Self {
        Builder
    }
}

/// A persistent state of a minimal master problem.
///
/// The implementation is trivial: this is just the master problem itself.
#[cfg_attr(feature = "serialize", derive(Serialize, Deserialize))]
#[derive(Debug, Clone)]
pub struct MinimalMasterState<Mn: Minorant>(MinimalMaster<Mn>);

impl<Mn: Minorant> Saveable for MinimalMaster<Mn> {
    type State = MinimalMasterState<Mn>;

    type Err = Infallible;

    fn get_state(&self) -> Result<Self::State, Infallible> {
        Ok(MinimalMasterState(self.clone()))
    }

    fn set_state(&mut self, state: MinimalMasterState<Mn>) -> Result<(), Infallible> {
        *self = state.0;
        Ok(())
    }
}

impl<Mn: Minorant> super::Builder<Mn> for Builder {
    type MasterProblem = MinimalMaster<Mn>;

    fn build(&mut self) -> Result<MinimalMaster<Mn>, MinimalMasterError> {
        MinimalMaster::new()
    }
}
