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

//! A wrapper around master problems to handle primal information.

use super::MasterProblem;
use crate::problem::SubgradientExtender;
use crate::{Aggregatable, Minorant, Real};

use either::Either;
use std::collections::HashMap;
use std::ops::{Deref, DerefMut};

/// A wrapper around `MasterProblem` to handle primal information.
///
/// For each minorant there can be an associated (generating) primal data.
/// Typically this data arises from a Lagrangian Relaxation approach. The primal
/// data is aggregated and handled along its associated minorant. However, the
/// primal data is irrelevant for the master problem, hence it is handled in
/// this wrapper.
pub struct PrimalMaster<M, P>
where
    M: MasterProblem,
{
    master: M,
    primals: HashMap<M::MinorantIndex, P>,
}

impl<M, P> PrimalMaster<M, P>
where
    M: MasterProblem,
    M::MinorantIndex: std::hash::Hash,
    P: Aggregatable,
{
    pub fn new(master: M) -> Self {
        PrimalMaster {
            master,
            primals: HashMap::new(),
        }
    }

    pub fn add_vars<E>(
        &mut self,
        vars: Vec<(Option<usize>, Real, Real)>,
        sgext: &mut SubgradientExtender<P, E>,
    ) -> Result<(), Either<M::Err, E>> {
        let primals = &self.primals;
        self.master.add_vars(&vars, |fidx, minidx, vars: &[usize]| {
            sgext(
                fidx,
                primals
                    .get(&minidx)
                    .expect("Extension for non-existing primal requested"),
                vars,
            )
        })
    }

    pub fn add_minorant(&mut self, fidx: usize, minorant: Minorant, primal: P) -> Result<(), M::Err> {
        let index = self.master.add_minorant(fidx, minorant)?;
        let old = self.primals.insert(index, primal);
        debug_assert!(old.is_none(), "Minorant index already used");
        Ok(())
    }

    /// Ensures that the aggregated submodels have the correct minorants.
    ///
    /// Subproblem that have too few minorants are extended using the latest
    /// aggregated minorant.
    pub fn fill_models(&mut self) -> Result<(), M::Err> {
        let primals = &mut self.primals;
        self.master.fill_models(|newindex, coeffs| {
            let newprimal = Aggregatable::combine(coeffs.map(|(i, alpha)| (alpha, &primals[&i])));
            let old = primals.insert(newindex, newprimal);
            debug_assert!(old.is_none(), "Minorant index already used");
        })
    }

    pub fn compress(&mut self) -> Result<(), M::Err> {
        let primals = &mut self.primals;
        self.master.compress(|newindex, coeffs| {
            let newprimal = Aggregatable::combine(
                coeffs.map(|(i, alpha)| (alpha, primals.remove(&i).expect("Missing primal to be compressed"))),
            );
            let old = primals.insert(newindex, newprimal);
            debug_assert!(old.is_none(), "Minorant index already used");
        })
    }

    /// Return the current aggregated primal information for a subproblem.
    ///
    /// The aggregated primal corresponds to the aggregated minorant of the last
    /// master solution.
    pub fn aggregated_primal(&self, fidx: usize) -> Result<P, M::Err> {
        Ok(P::combine(self.master.opt_multipliers(fidx).map(|(i, alpha)| {
            (alpha, self.primals.get(&i).expect("Missing primal to be aggregated"))
        })))
    }
}

impl<M, P> Deref for PrimalMaster<M, P>
where
    M: MasterProblem,
    P: Aggregatable,
{
    type Target = M;
    fn deref(&self) -> &M {
        &self.master
    }
}

impl<M, P> DerefMut for PrimalMaster<M, P>
where
    M: MasterProblem,
    P: Aggregatable,
{
    fn deref_mut(&mut self) -> &mut M {
        &mut self.master
    }
}
