/*
 * 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/>
 */

#![allow(non_upper_case_globals)]

//! Example implementation for a capacitated facility location problem.

use better_panic;
use log::{info, Level};
use rustop::opts;
use std::error::Error;
use std::fmt::Write;
use std::io::Write as _;
use std::sync::Arc;

use env_logger::{self, fmt::Color};
use ordered_float::NotNan;
use threadpool::ThreadPool;

use bundle::problem::{FirstOrderProblem as ParallelProblem, ResultSender};
use bundle::solver::sync::{DefaultSolver, NoBundleSolver};
use bundle::{dvec, DVector, Minorant, Real};

const Nfac: usize = 3;
const Ncus: usize = 5;
const F: [Real; Nfac] = [1000.0, 1000.0, 1000.0];
const CAP: [Real; Nfac] = [500.0, 500.0, 500.0];
const C: [[Real; Ncus]; Nfac] = [
    [4.0, 5.0, 6.0, 8.0, 10.0], //
    [6.0, 4.0, 3.0, 5.0, 8.0],  //
    [9.0, 7.0, 4.0, 3.0, 4.0],  //
];
const DEMAND: [Real; 5] = [80.0, 270.0, 250.0, 160.0, 180.0];

#[derive(Debug)]
enum EvalError {
    Customer(usize),
}

impl std::fmt::Display for EvalError {
    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            EvalError::Customer(i) => writeln!(fmt, "Customer subproblem {} failed", i),
        }
    }
}

impl Error for EvalError {}

struct CFLProblem {
    pool: Option<ThreadPool>,
}

impl CFLProblem {
    fn new() -> CFLProblem {
        CFLProblem { pool: None }
    }
}

fn solve_facility_problem(j: usize, lambda: &[Real]) -> Result<(Real, Minorant, DVector), EvalError> {
    // facility subproblem max\{-F[j] + CAP[j] * lambda[j] * y[j] | y_j \in \{0,1\}\}
    let objective;
    let mut subg = dvec![0.0; lambda.len()];
    let primal;
    if -F[j] + CAP[j] * lambda[j] > 0.0 {
        objective = -F[j] + CAP[j] * lambda[j];
        subg[j] = CAP[j];
        primal = dvec![1.0; 1];
    } else {
        objective = 0.0;
        primal = dvec![0.0; 1];
    }

    Ok((
        objective,
        Minorant {
            constant: objective,
            linear: subg,
        },
        primal,
    ))
}

fn solve_customer_problem(i: usize, lambda: &[Real]) -> Result<(Real, Minorant, DVector), EvalError> {
    // customer subproblem
    let opt = (0..Nfac)
        .map(|j| (j, -DEMAND[i] * C[j][i] - DEMAND[i] * lambda[j]))
        .flat_map(|x| NotNan::new(x.1).map(|y| (x.0, y)))
        .max_by_key(|x| x.1)
        .ok_or(EvalError::Customer(i))?;

    let objective = opt.1.into_inner();
    let mut subg = dvec![0.0; lambda.len()];
    subg[opt.0] = -DEMAND[i];
    let mut primal = dvec![0.0; Nfac];
    primal[opt.0] = 1.0;

    Ok((
        objective,
        Minorant {
            constant: objective,
            linear: subg,
        },
        primal,
    ))
}

impl ParallelProblem for CFLProblem {
    type Err = EvalError;

    type Primal = DVector;

    fn num_variables(&self) -> usize {
        Nfac
    }

    fn lower_bounds(&self) -> Option<Vec<Real>> {
        Some(vec![0.0; ParallelProblem::num_variables(self)])
    }

    fn num_subproblems(&self) -> usize {
        Nfac + Ncus
    }

    fn start(&mut self) {
        let pool = ThreadPool::new(4);
        self.pool = Some(pool)
    }

    fn stop(&mut self) {
        self.pool.take();
    }

    fn evaluate<S>(&mut self, i: usize, y: Arc<DVector>, tx: S) -> Result<(), Self::Err>
    where
        S: ResultSender<Self> + 'static,
    {
        if self.pool.is_none() {
            self.start()
        }
        let y = y.clone();
        self.pool.as_ref().unwrap().execute(move || {
            let res = if i < Nfac {
                solve_facility_problem(i, &y)
            } else {
                solve_customer_problem(i - Nfac, &y)
            };
            match res {
                Ok((value, minorant, primal)) => {
                    if tx.objective(value).is_err() {
                        return;
                    }
                    if tx.minorant(minorant, primal).is_err() {
                        return;
                    }
                }
                Err(err) => if tx.error(err).is_err() {},
            };
        });
        Ok(())
    }
}

fn show_primals<F>(f: F) -> Result<(), Box<dyn Error>>
where
    F: Fn(usize) -> DVector,
{
    let mut totalobj = 0.0;
    for i in 0..Ncus {
        let x = f(Nfac + i);
        let mut obj = 0.0;
        let mut s = String::new();
        write!(s, "x[{}] =", i)?;
        for j in 0..Nfac {
            write!(s, " {:.2}", x[j])?;
            obj += DEMAND[i] * C[j][i] * x[j];
        }
        write!(s, "  objval = {:.2}", obj)?;
        info!("{}", s);
        totalobj += obj;
    }

    let mut obj = 0.0;
    let mut s = String::new();
    write!(s, "y =")?;
    for j in 0..Nfac {
        let y = f(j)[0];
        write!(s, " {:.2}", y)?;
        obj += F[j] * y;
    }
    totalobj += obj;
    write!(s, "  objval = {:.2}", obj)?;
    info!("{}", s);
    info!("total objective = {:.2}", totalobj);

    Ok(())
}

fn main() -> Result<(), Box<dyn Error>> {
    better_panic::install();
    env_logger::builder()
        .format(|buf, record| {
            let mut style = buf.style();
            let color = match record.level() {
                Level::Error | Level::Warn => Color::Red,
                Level::Trace => Color::Blue,
                Level::Debug => Color::Yellow,
                _ => Color::White,
            };
            style.set_color(color);
            writeln!(
                buf,
                "{}{:5}{} {}",
                style.value("["),
                style.value(record.level()),
                style.value("]"),
                style.value(record.args())
            )
        })
        .init();

    let (args, _) = opts! {
        synopsis "Solver a simple capacitated facility location problem";
        opt minimal:bool, desc:"Use the minimal master model";
    }
    .parse_or_exit();
    if !args.minimal {
        let mut slv = DefaultSolver::<_>::new(CFLProblem::new());
        slv.terminator.termination_precision = 1e-9;
        slv.master.max_bundle_size = 5;
        slv.solve()?;

        show_primals(|i| slv.aggregated_primal(i).unwrap())?;
    } else {
        let mut slv = NoBundleSolver::<_>::new(CFLProblem::new());
        slv.terminator.termination_precision = 1e-5;
        slv.solve()?;

        show_primals(|i| slv.aggregated_primal(i).unwrap())?;
    }

    Ok(())
}
