/*
 * 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 better_panic;
use bundle::dvec;
use env_logger;
use env_logger::fmt::Color;
use log::{debug, info, Level};
use rustop::opts;
use std::convert::Infallible;
use std::io::Write;
use std::sync::Arc;
use std::thread;

use bundle::master::{Builder as MasterBuilder, FullMasterBuilder, MinimalMasterBuilder};
use bundle::problem::{FirstOrderProblem as ParallelProblem, ResultSender};
use bundle::saveable::Saveable;
use bundle::solver::sync::{Error, Solver};
use bundle::terminator::StandardTerminator;
use bundle::weighter::HKWeighter;
use bundle::{DVector, Real};

#[derive(Clone)]
struct QuadraticProblem {
    a: [[Real; 2]; 2],
    b: [Real; 2],
    c: Real,
}

impl QuadraticProblem {
    fn new() -> QuadraticProblem {
        QuadraticProblem {
            a: [[5.0, 1.0], [1.0, 4.0]],
            b: [-12.0, -10.0],
            c: 3.0,
        }
    }
}

impl ParallelProblem for QuadraticProblem {
    // No error can occur.
    //
    // Even better would by `type Err = !;` which is only available
    // with `#![feature(never_type)]` as of Rust 1.60.
    type Err = Infallible;

    type Minorant = (Real, DVector);

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

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

    fn start(&mut self) {}

    fn stop(&mut self) {}

    fn evaluate<S>(&mut self, fidx: usize, x: Arc<DVector>, tx: S) -> Result<(), Self::Err>
    where
        S: ResultSender<Self> + 'static,
    {
        let x = x.clone();
        let p = self.clone();
        thread::spawn(move || {
            assert_eq!(fidx, 0);
            let mut objective = p.c;
            let mut g = dvec![0.0; 2];

            for i in 0..2 {
                g[i] += (0..2).map(|j| p.a[i][j] * x[j]).sum::<Real>();
                objective += x[i] * (g[i] + p.b[i]);
                g[i] = 2.0 * g[i] + p.b[i];
            }

            debug!("Evaluation at {:?}", x);
            debug!("  objective={}", objective);
            debug!("  subgradient={}", g);

            tx.objective(objective).unwrap();
            tx.minorant((objective, g)).unwrap();
        });
        Ok(())
    }
}

fn run<M>() -> Result<(), Box<dyn std::error::Error>>
where
    M: MasterBuilder<(Real, DVector)> + Default,
    M::MasterProblem: Saveable,
{
    {
        info!("-- Solve completely --");
        let f = QuadraticProblem::new();
        let mut solver = Solver::<_, StandardTerminator, HKWeighter, M>::new(f);
        solver.solve()?;
    }

    {
        info!("-- Solve with interrupt and resume --");
        let f = QuadraticProblem::new();
        let mut solver = Solver::<_, StandardTerminator, HKWeighter, M>::new(f);

        match solver.solve_with_limit(3) {
            Ok(_) | Err(Error::IterationLimit { .. }) => (),
            Err(err) => return Err(err.into()),
        }

        info!("- interrupt -");
        solver.solve_iter(1000)?;
    }

    {
        info!("-- Solve with interrupt to state and resume --");
        let f = QuadraticProblem::new();
        let mut solver = Solver::<_, StandardTerminator, HKWeighter, M>::new(f);

        match solver.solve_with_limit(3) {
            Ok(_) | Err(Error::IterationLimit { .. }) => (),
            Err(err) => return Err(err.into()),
        }
        let s = solver.get_state()?;

        info!("- interrupt -");
        let f = solver.into_problem();
        let mut solver2 = Solver::<_, StandardTerminator, HKWeighter, M>::new(f);
        solver2.set_state(s)?;
        solver2.solve_iter(1000)?;
    }

    Ok(())
}

fn main() -> Result<(), Box<dyn std::error::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 quadratic optimization problem";
        opt minimal:bool, desc:"Use the minimal master model";
    }
    .parse_or_exit();

    if !args.minimal {
        info!("Use full master problem");
        run::<FullMasterBuilder>()?;
    } else {
        info!("Use minimal master problem");
        run::<MinimalMasterBuilder>()?;
    }

    Ok(())
}
