// Copyright (c) 2016, 2017 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 mcf;
use {DVector, FirstOrderProblem, Minorant, Real, SimpleEvaluation};

use std::f64::INFINITY;
use std::fs::File;
use std::io::Read;
use std::result::Result;

use failure::Error;

/// A solver error.
#[derive(Debug, Fail)]
#[fail(display = "Format error: {}", msg)]
pub struct MCFFormatError {
    msg: String,
}

#[derive(Clone, Copy, Debug)]
struct ArcInfo {
    arc: usize,
    src: usize,
    snk: usize,
}

#[derive(Clone, Copy, Debug)]
struct Elem {
    ind: usize,
    val: Real,
}

pub struct MMCFProblem {
    pub multimodel: bool,

    nets: Vec<mcf::Solver>,
    lhs: Vec<Vec<Vec<Elem>>>,
    rhs: DVector,
    rhsval: Real,
    cbase: Vec<DVector>,
    c: Vec<DVector>,
}

impl MMCFProblem {
    pub fn read_mnetgen(basename: &str) -> Result<MMCFProblem, Error> {
        let mut buffer = String::new();
        {
            let mut f = try!(File::open(&format!("{}.nod", basename)));
            try!(f.read_to_string(&mut buffer));
        }
        let fnod = buffer
            .split_whitespace()
            .map(|x| x.parse::<usize>().unwrap())
            .collect::<Vec<_>>();

        if fnod.len() != 4 {
            return Err(MCFFormatError {
                msg: format!("Expected 4 numbers in {}.nod, but got {}", basename, fnod.len()),
            }.into());
        }

        let ncom = fnod[0];
        let nnodes = fnod[1];
        let narcs = fnod[2];
        let ncaps = fnod[3];

        // read nodes
        let mut nets = Vec::with_capacity(ncom);
        for _ in 0..ncom {
            nets.push(try!(mcf::Solver::new(nnodes)))
        }
        {
            let mut f = try!(File::open(&format!("{}.sup", basename)));
            buffer.clear();
            try!(f.read_to_string(&mut buffer));
        }
        for line in buffer.lines() {
            let mut data = line.split_whitespace();
            let node = try!(data.next().unwrap().parse::<usize>());
            let com = try!(data.next().unwrap().parse::<usize>());
            let supply = try!(data.next().unwrap().parse::<Real>());
            try!(nets[com - 1].set_balance(node - 1, supply));
        }

        // read arcs
        let mut arcmap = vec![vec![]; ncom];
        let mut cbase = vec![dvec![]; ncom];

        // lhs nonzeros
        let mut lhsidx = vec![vec![vec![]; ncom]; ncaps];
        {
            let mut f = try!(File::open(&format!("{}.arc", basename)));
            buffer.clear();
            try!(f.read_to_string(&mut buffer));
        }
        for line in buffer.lines() {
            let mut data = line.split_whitespace();
            let arc = try!(data.next().unwrap().parse::<usize>()) - 1;
            let src = try!(data.next().unwrap().parse::<usize>()) - 1;
            let snk = try!(data.next().unwrap().parse::<usize>()) - 1;
            let com = try!(data.next().unwrap().parse::<usize>()) - 1;
            let cost = try!(data.next().unwrap().parse::<Real>());
            let cap = try!(data.next().unwrap().parse::<Real>());
            let mt = try!(data.next().unwrap().parse::<isize>()) - 1;
            assert!(
                arc < narcs,
                format!("Wrong arc number (got: {}, expected in 1..{})", arc + 1, narcs)
            );
            // set internal coeff
            let coeff = arcmap[com].len();
            arcmap[com].push(ArcInfo {
                arc: arc + 1,
                src: src + 1,
                snk: snk + 1,
            });
            // add arc
            try!(nets[com].add_arc(src, snk, cost, if cap < 0.0 { INFINITY } else { cap }));
            // set objective
            cbase[com].push(cost); // + 1e-6 * coeff
                                   // add to mutual capacity constraint
            if mt >= 0 {
                lhsidx[mt as usize][com].push(coeff);
            }
        }

        // read rhs of coupling constraints
        {
            let mut f = try!(File::open(&format!("{}.mut", basename)));
            buffer.clear();
            try!(f.read_to_string(&mut buffer));
        }
        let mut rhs = dvec![0.0; ncaps];
        for line in buffer.lines() {
            let mut data = line.split_whitespace();
            let mt = try!(data.next().unwrap().parse::<usize>()) - 1;
            let cap = try!(data.next().unwrap().parse::<Real>());
            rhs[mt] = cap;
        }

        // set lhs
        let mut lhs = vec![vec![vec![]; ncom]; ncaps];
        for i in 0..ncaps {
            for fidx in 0..ncom {
                lhs[i][fidx] = lhsidx[i][fidx].iter().map(|&j| Elem { ind: j, val: 1.0 }).collect();
            }
        }

        Ok(MMCFProblem {
            multimodel: false,
            nets,
            lhs,
            rhs,
            rhsval: 0.0,
            cbase,
            c: vec![dvec![]; ncom],
        })
    }

    /// Compute costs for a primal solution.
    pub fn get_primal_costs(&self, fidx: usize, primals: &[DVector]) -> Real {
        if self.multimodel {
            primals[0]
                .iter()
                .enumerate()
                .map(|(i, x)| x * self.cbase[fidx][i])
                .sum()
        } else {
            let mut sum = 0.0;
            for (fidx, p) in primals.iter().enumerate() {
                for (i, x) in p.iter().enumerate() {
                    sum += x * self.cbase[fidx][i];
                }
            }
            sum
        }
    }

    /// Aggregate primal vectors.
    pub fn aggregate_primals_ref(&self, primals: &[(Real, &Vec<DVector>)]) -> Vec<DVector> {
        let mut aggr = primals[0]
            .1
            .iter()
            .map(|x| {
                let mut r = dvec![];
                r.scal(primals[0].0, x);
                r
            })
            .collect::<Vec<_>>();

        for &(alpha, primal) in &primals[1..] {
            for (j, x) in primal.iter().enumerate() {
                aggr[j].add_scaled(alpha, x);
            }
        }

        aggr
    }
}

impl<'a> FirstOrderProblem<'a> for MMCFProblem {
    type Primal = Vec<DVector>;

    type EvalResult = SimpleEvaluation<Vec<DVector>>;

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

    fn lower_bounds(&self) -> Option<Vec<Real>> {
        Some(vec![0.0; self.lhs.len()])
    }

    fn upper_bounds(&self) -> Option<Vec<Real>> {
        None
    }

    fn num_subproblems(&self) -> usize {
        if self.multimodel {
            self.nets.len()
        } else {
            1
        }
    }

    #[allow(unused_variables)]
    fn evaluate(
        &'a mut self,
        fidx: usize,
        y: &[Real],
        nullstep_bound: Real,
        relprec: Real,
    ) -> Result<Self::EvalResult, Error> {
        // compute costs
        self.rhsval = 0.0;
        for i in 0..self.c.len() {
            self.c[i].clear();
            self.c[i].extend(self.cbase[i].iter());
        }

        for (i, &y) in y.iter().enumerate().filter(|&(i, &y)| y != 0.0) {
            self.rhsval += self.rhs[i] * y;
            for (fidx, c) in self.c.iter_mut().enumerate() {
                for elem in &self.lhs[i][fidx] {
                    c[elem.ind] += y * elem.val;
                }
            }
        }

        debug!("y={:?}", y);
        for i in 0..self.nets.len() {
            debug!("c[{}]={}", i, self.c[i]);
            try!(self.nets[i].set_objective(&self.c[i]));
        }

        // solve subproblems
        for (i, net) in self.nets.iter_mut().enumerate() {
            try!(net.solve());
            debug!("c[{}]={}", i, try!(net.objective()));
        }

        // compute minorant
        if self.multimodel {
            let objective;
            let mut subg;
            if fidx == 0 {
                subg = self.rhs.clone();
                objective = self.rhsval - try!(self.nets[fidx].objective());
            } else {
                subg = dvec![0.0; self.rhs.len()];
                objective = -try!(self.nets[fidx].objective());
            }

            let sol = try!(self.nets[fidx].get_solution());
            for (i, lhs) in self.lhs.iter().enumerate() {
                subg[i] -= lhs[fidx].iter().map(|elem| elem.val * sol[elem.ind]).sum::<Real>();
            }

            Ok(SimpleEvaluation {
                objective,
                minorants: vec![(
                    Minorant {
                        constant: objective,
                        linear: subg,
                    },
                    vec![sol],
                )],
            })
        } else {
            let mut objective = self.rhsval;
            let mut sols = Vec::with_capacity(self.nets.len());
            for i in 0..self.nets.len() {
                objective -= try!(self.nets[i].objective());
                sols.push(try!(self.nets[i].get_solution()));
            }

            let mut subg = self.rhs.clone();
            for (i, lhs) in self.lhs.iter().enumerate() {
                for (fidx, flhs) in lhs.iter().enumerate() {
                    subg[i] -= flhs.iter().map(|elem| elem.val * sols[fidx][elem.ind]).sum::<Real>();
                }
            }

            Ok(SimpleEvaluation {
                objective,
                minorants: vec![(
                    Minorant {
                        constant: objective,
                        linear: subg,
                    },
                    sols,
                )],
            })
        }
    }

    fn aggregate_primals(&mut self, primals: Vec<(Real, Vec<DVector>)>) -> Vec<DVector> {
        self.aggregate_primals_ref(&primals.iter().map(|&(alpha, ref x)| (alpha, x)).collect::<Vec<_>>())
    }
}
