/*
 * 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 crossbeam::channel::unbounded as channel;
use log::{info, Level};
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::parallel::{EvalResult, FirstOrderProblem as ParallelProblem, ParallelSolver, ResultSender};
use bundle::{dvec, DVector, Minorant, Real};
use bundle::{DefaultSolver, FirstOrderProblem, SimpleEvaluation, StandardTerminator};

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),
    NoObjective(usize),
}

impl std::fmt::Display for EvalError {
    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
        match self {
            EvalError::Customer(i) => writeln!(fmt, "Customer subproblem {} failed", i),
            EvalError::NoObjective(i) => writeln!(fmt, "No objective value generated for subproblem {}", 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 FirstOrderProblem for CFLProblem {
    type Err = EvalError;

    type Primal = DVector;

    type EvalResult = SimpleEvaluation<Self::Primal>;

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

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

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

    fn evaluate(
        &mut self,
        index: usize,
        lambda: &[Real],
        _nullstep_bound: Real,
        _relprec: Real,
    ) -> Result<Self::EvalResult, Self::Err> {
        let (tx, rx) = channel();
        ParallelProblem::evaluate(self, index, Arc::new(lambda.iter().cloned().collect()), index, tx)?;
        let mut objective = None;
        let mut minorants = vec![];

        for r in rx {
            match r {
                Ok(EvalResult::ObjectiveValue { value, .. }) => objective = Some(value),
                Ok(EvalResult::Minorant { minorant, primal, .. }) => minorants.push((minorant, primal)),
                _ => break,
            }
        }

        Ok(SimpleEvaluation {
            objective: objective.ok_or(EvalError::NoObjective(index))?,
            minorants,
        })
    }
}

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<I>(
        &mut self,
        i: usize,
        y: Arc<DVector>,
        index: I,
        tx: ResultSender<I, Self::Primal, Self::Err>,
    ) -> Result<(), Self::Err>
    where
        I: Send + Copy + '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.send(Ok(EvalResult::ObjectiveValue { index, value })).is_err() {
                        return;
                    }
                    if tx
                        .send(Ok(EvalResult::Minorant {
                            index,
                            minorant,
                            primal,
                        }))
                        .is_err()
                    {
                        return;
                    }
                }
                Err(err) => if tx.send(Err(err)).is_err() {},
            };
        });
        Ok(())
    }
}

fn main() -> Result<(), Box<Error>> {
    // env_logger::builder()
    //     .default_format_timestamp(false)
    //     .default_format_module_path(false)
    //     .init();
    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 mut slv = DefaultSolver::new(CFLProblem::new())?;
        slv.terminator = Box::new(StandardTerminator {
            termination_precision: 1e-9,
        });
        slv.solve()?;

        for i in 0..Ncus {
            let x = slv.aggregated_primals(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);
        }

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

    {
        let mut slv = ParallelSolver::<_>::new(CFLProblem::new())?;
        slv.solve()?;
    }

    Ok(())
}
