//! Create and manipulate semi-definite problems.

extern crate ndarray;
extern crate num;
extern crate sprs;

use self::num::*;
use self::sprs::{CsMat, TriMat};
use crate::algebra::*;
use crate::flag::Flag;
use crate::operator::*;

use std::fmt::Display;
use std::fs::File;
use std::io::*;
use std::io::{BufWriter, Write};
use std::path::PathBuf;
use std::str::FromStr;

/// The optimisation problems over flags are translated into a
/// sdp problem in the sdpa format.
///
/// Shape of the matrices:
///
/// For each i in ineqs (where i is itself a vector of inequalities):
///    A diagonal block of size i.len()
///
/// For each cs:
///    A block with the size od cs.input_matrix

/// An optimization problem expressed in flags algebra.
#[derive(Debug)]
pub struct Problem<N, F> {
    /// Set of contraint inequalities.
    pub ineqs: Vec<Ineq<N, F>>,
    /// Set of Cauchy-Schwarz inequlities to be used.
    pub cs: Vec<MulAndUnlabeling<F>>,
    /// Vector to be optimized.
    pub obj: QFlag<N, F>,
}

fn write_coeff<N>(
    file: &mut BufWriter<File>,
    mat_num: usize,
    block_num: usize,
    i: usize,
    j: usize,
    value: N,
) -> Result<()>
where
    N: Display + Zero + PartialEq + Copy,
{
    if value != N::zero() {
        writeln!(
            file,
            "{} {} {} {} {}",
            mat_num,
            block_num + 1,
            i + 1,
            j + 1,
            value
        )?;
    }
    Ok(())
}

impl<N, F> Problem<N, F>
where
    N: Display + Zero + Copy + PartialEq,
    F: Flag,
{
    /// Panic if the size of the basis involved are inconsistent.
    pub fn check(&self) {
        let b = self.obj.basis;
        for ineq in &self.ineqs {
            assert_eq!(ineq.meta.basis, b);
        }
        for cs in &self.cs {
            assert_eq!(cs.output_basis(), b);
        }
    }

    /// Write a description of the problem as a comment in the .sdpa file
    fn write_header(&self, file: &mut BufWriter<File>) -> Result<()> {
        writeln!(file, "* sdp problem generated by Rust")?;
        writeln!(file, "* Flag: {}", F::NAME)?;
        writeln!(file, "* Basis: {:} ({} flags)", self.obj.basis, self.obj.basis.get().len())?;
        writeln!(file, "* {} groups of linear constraints", self.ineqs.len())?;
        writeln!(file, "* {} Cauchy-Schwarz constraints", self.cs.len())?;
        writeln!(
            file,
            "* Minimizing {} (scale {})",
            self.obj.expr, self.obj.scale
        )?;
        for i in &self.ineqs {
            writeln!(
                file,
                "* {}\t≥ {} ({})",
                i.meta.flag_expr,
                i.meta.bound_expr,
                i.data.len()
            )?;
        }
        for cs in &self.cs {
            writeln!(file,
                     "* Cauchy-Schwarz: {}x{}; {:?}",
                     cs.split.left_basis().size,
                     cs.split.right_basis().size,
                     cs.split.right_basis().t,
            )?;
        }
        writeln!(file, "*")
    }
    /// Write the semi-definite program in the file `filename` in the sdpa format.
    pub fn write_sdpa(&self, filename: &str) -> Result<()> {
        self.check();
        let mut filename = PathBuf::from(filename);
        let _ = filename.set_extension("sdpa");
        let mut file = BufWriter::new(File::create(&filename)?);
        self.write_header(&mut file)?;
        let cs_mat: Vec<_> = self.cs.iter().map(|x| x.get()).collect();
        // Line 1: Number of constraints = size of the basis
        writeln!(file, "{}", self.obj.data.len())?;
        // Line 2: Number of blocks (one for each constraint)
        writeln!(file, "{}", self.ineqs.len() + self.cs.len())?;
        // Line 3: Sizes of the blocks
        for ineq in &self.ineqs {
            write!(file, "-{} ", ineq.data.len())?;
        }
        for split in &cs_mat {
            write!(file, "{} ", split[0].rows())?;
        }
        writeln!(file)?;
        // Line 4: vector ai
        for v in &self.obj.data {
            write!(file, "{} ", v)?;
        }
        writeln!(file)?;
        // Lines 5+: body
        // Matrix 0: Objective
        for (block_num, ineq) in self.ineqs.iter().enumerate() {
            for (i, inequality) in ineq.data.iter().enumerate() {
                write_coeff(&mut file, 0, block_num, i, i, inequality.bound)?;
            }
        }
        writeln!(file)?;
        // Matrices 1+:
        // Inequaltity blocks
        for (block_num, ineq) in self.ineqs.iter().enumerate() {
            for (i, ineqdata) in ineq.data.iter().enumerate() {
                for (mat_num, &v) in ineqdata.flag.iter().enumerate() {
                    write_coeff(&mut file, mat_num + 1, block_num, i, i, v)?;
                }
            }
        }
        writeln!(file)?;
        // Cs blocks
        let offset = self.ineqs.len();
        for (block_num, line) in cs_mat.iter().enumerate() {
            for (mat_num, matrix) in line.iter().enumerate() {
                for (&v, (i, j)) in matrix.iter() {
                    if i <= j {
                        write_coeff(&mut file, mat_num + 1, block_num + offset, i, j, v)?;
                    }
                }
            }
        }
        Ok(())
    }
}

// A line in the .sdpa format
struct SdpaCoeff {
    mat: usize,
    block: usize,
    i: usize,
    j: usize,
    val: f64,
}

impl FromStr for SdpaCoeff {
    type Err = Box<dyn std::error::Error>;

    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        let mut iter = s.split_whitespace();
        let mat = iter.next().unwrap().parse::<usize>()?;
        let block = iter.next().unwrap().parse::<usize>()?;
        let i = iter.next().unwrap().parse::<usize>()?;
        let j = iter.next().unwrap().parse::<usize>()?;
        let val = iter.next().unwrap().parse::<f64>()?;
        assert_eq!(iter.next(), None);
        Ok(SdpaCoeff {
            mat,
            block,
            i,
            j,
            val,
        })
    }
}

// ===== Metadata =====

#[derive(Debug, Clone)]
pub enum BlockInfo<F> {
    Inequality(IneqMeta<F>),
    CauchySchwarz(MulAndUnlabeling<F>),
}

#[derive(Debug, Clone)]
pub struct BlockMeta<F> {
    size: usize,
    diagonal: bool,
    info: BlockInfo<F>,
}

// ========== Certificate ==========
#[derive(Debug, Clone)]
pub struct Certificate<N> {
    pub y: Vec<N>,
    pub z: Vec<CsMat<N>>, // 1
    pub x: Vec<CsMat<N>>, // 2
}

impl Certificate<f64> {
    pub fn from_file<N, F: Flag>(pb: &Problem<N, F>, name: &str) -> Result<Self> {
        let file = File::open(name)?;
        let mut buf = BufReader::new(file).lines();
        // The value of the vector y is on the first line
        let y = buf
            .next()
            .unwrap()?
            .split_whitespace()
            .map(|x| x.parse::<f64>().unwrap())
            .collect();
        // Prepare the space for the matrices z and x
        let mut tri_mat = [Vec::new(), Vec::new()];
        for i in &pb.ineqs {
            let len = i.data.len();
            for m in &mut tri_mat {
                m.push(TriMat::new((len, len)))
            }
        }
        for cs in &pb.cs {
            let n = cs.get()[0].rows();
            for m in &mut tri_mat {
                m.push(TriMat::new((n, n)))
            }
        }
        for line in buf {
            let l = line.unwrap().parse::<SdpaCoeff>().unwrap();
            tri_mat[l.mat - 1][l.block - 1].add_triplet(l.i - 1, l.j - 1, l.val);
            // If the value is not on the diagonal, add the symetric coeff
            if l.i != l.j {
                tri_mat[l.mat - 1][l.block - 1].add_triplet(l.j - 1, l.i - 1, l.val);
            }
        }
        // convert the triplet matrices to sparse matrices
        let mut sprs_mat = tri_mat
            .into_iter()
            .map(|mat| mat.iter().map(|block| block.to_csc()).collect());
        Ok(Self {
            y,
            z: sprs_mat.next().unwrap(), // First matrix: z
            x: sprs_mat.next().unwrap(), // Second matrix: x
        })
    }
    pub fn value<F, N>(&self, pb: &Problem<N, F>) -> f64
    where
        N: Clone + ToPrimitive,
        F: Flag,
    {
        let mut res = 0.;
        // For each block of inequalities
        for (block, ineqs) in pb.ineqs.iter().enumerate() {
            for (v, (i, j)) in self.x[block].iter() {
                assert_eq!(i, j);
                let bound: f64 = NumCast::from(ineqs.data[i].bound.clone()).unwrap();
                res += v * bound;
            }
        }
        res
    }
    pub fn to_vec<F>(&self, b: Basis<F>, threshold: f64) -> QFlag<f64, F>
    where
        F: Flag,
    {
        b.from_vec(
            self.y
                .iter()
                .map(|x| if x.abs() < threshold { 0. } else { *x })
                .collect(),
        )
    }
    pub fn with_threshold(mut self, threshold: f64) -> Self {
        for x in &mut self.y {
            if x.abs() < threshold {
                *x = 0.
            }
        }
        self
    }

    pub fn print_info(&self) {
        let zero_y = self.y.iter().filter(|&&x| x == 0.).count();
        println!("y: size {} with {} zeroes", self.y.len(), zero_y);
        for (i, m) in self.x.iter().enumerate() {
            println!(
                "x {}: size {} with {} non-zeroes {}",
                i,
                m.outer_dims(),
                m.iter().filter(|&(&x, _)| { x >= 0.0001 }).count(),
                if m.data().len() >= 2 {
                    m.data()[0] - m.data()[1]
                } else {
                    42.
                }
            );
        }
        // pub z: Vec<CsMat<N>>,
        //  pub x: Vec<CsMat<N>>,
    }
}