use mtxlib::*;
use krysys::{*, minres::Minres};

fn sierpinski_matrix(nrows: usize, ncols: usize) -> (Vec<usize>, Vec<usize>, Vec<f64>) {
    let mut rows: Vec<usize> = Vec::new();
    let mut cols: Vec<usize> = Vec::new();
    let mut vals: Vec<f64> = Vec::new();

    println!("Sparsity pattern of constraint coefficient matrix:");

    for r in 0..nrows {
        for c in 0..ncols {
            if (r & c) == 0 {
                rows.push(r);
                cols.push(c);
                vals.push(1.0);
                print!("* ");
            } else {
                print!("  ");
            }

            if c == (ncols - 1) {
                println!();
            }
        }
    }
    (rows, cols, vals)
}

struct QuadraticProgramBlackBox<'a, T> {
    // constraint coefficient matrix
    coeff: st_ops::GenericMatrix<'a, T>,

    // problem dimensions
    var_sz: usize,
    constraint_sz: usize,
}

impl<'a, T: Float, H> LinearMap<T, H> for QuadraticProgramBlackBox<'a, T> {
    fn apply(&self, product: &mut [T], input: &[T], _helper: &mut H) {
        product.reset();

        // First part of product
        self.coeff.lmv_mul(&mut product[0..self.var_sz], &input[self.var_sz..(self.var_sz + self.constraint_sz)]);

        product[0..self.var_sz].iter_mut().zip(input[0..self.var_sz].iter()).for_each(|(p, i)| *p += *i);

        // Second part of product
        self.coeff.mv_mul(&mut product[self.var_sz..(self.var_sz + self.constraint_sz)], input);
    }
}

fn main() {
    // Minimize
    //      1/2 x^T x
    // subject to
    //      A x = 1
    // where A is a wide Sierpinski matrix

    let constraint_sz: usize = 2_i32.pow(5) as usize;
    let var_sz: usize = 2 * constraint_sz;

    // Specify tolerance
    let tol: f64 = 1e-8;

    // Specify max iterations allowed
    let max_iters = 10_000;

    // Select an appropriate plugin.
    let mut do_nothing_plugin = DoNothing;
    let mut logger = ConsoleLogger::extend(&mut do_nothing_plugin);
    let mut plugin = StopWatch::extend(&mut logger);

    let opts =
        minres::Settings { abs_tolerance: tol, rel_tolerance: tol, max_iterations: max_iters, ..Default::default() };

    // System of equality constraints: coefficient matrix has a Sierpinski triangle pattern
    let (rows, cols, vals) = sierpinski_matrix(constraint_sz, var_sz);

    let total_nnz = vals.len();

    let mat = st_ops::GenericMatrix::new(&rows, &cols, &vals);

    let mut constraint_rhs: Vec<f64> = vec![0.0; constraint_sz];

    mat.mv_mul(&mut constraint_rhs, &vec![1.0; var_sz]);
    constraint_rhs[0] = 0.0;

    let black_box = QuadraticProgramBlackBox { coeff: mat, var_sz, constraint_sz };

    let mut sol: Vec<f64> = vec![0.0; var_sz + constraint_sz];
    let mut b: Vec<f64> = vec![0.0; var_sz];

    b.append(&mut constraint_rhs.clone());

    let mut aux1: Vec<f64> = vec![0.0; var_sz + constraint_sz];
    let mut aux2: Vec<f64> = vec![0.0; var_sz + constraint_sz];
    let mut aux3: Vec<f64> = vec![0.0; var_sz + constraint_sz];
    let mut aux4: Vec<f64> = vec![0.0; var_sz + constraint_sz];
    let mut aux5: Vec<f64> = vec![0.0; var_sz + constraint_sz];

    let minres::Report { status, residual, iterations: iters } = black_box.minres(
        &mut sol,
        &b,
        &opts,
        &mut plugin,
        minres::Buckets::new(&mut aux1, &mut aux2, &mut aux3, &mut aux4, &mut aux5),
    );

    println!("Primal Solution:");
    for val in sol[0..var_sz].iter() {
        println!("  {val: <+13.7e}");
    }

    println!(" ");
    println!("Dual Solution:");
    for val in sol[var_sz..sol.len()].iter() {
        println!("  {val: <+13.7e}");
    }

    println!("Number of variables: {var_sz}");
    println!("Number of constraints: {constraint_sz}");

    let kkt_sz = var_sz + constraint_sz;
    let kkt_nnz = var_sz + 2 * total_nnz;

    println!("Size of KKT matrix: {}", kkt_sz);
    println!("Number of KKT matrix non-zeros: {}", kkt_nnz);

    println!("Total iterations: {iters}");
    println!("Residual achieved: {residual:.7e}");

    let total_flop = (kkt_nnz + 11 * kkt_sz) * iters + kkt_nnz + 5 * kkt_sz;
    println!("Total floating point operations: {}", total_flop);
    println!("Elapsed time: {:.7e} s.", plugin.elapsed_time_in_seconds());

    println!("Status: {status:?}")
}
