// Copyright (c) 2016, 2017, 2018 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/>
//

//! Finite-dimensional sparse and dense vectors.

use crate::Real;
use std::fmt;
use std::ops::{Deref, DerefMut};
// use std::cmp::min;
use std::iter::FromIterator;
use std::vec::IntoIter;

/// Type of dense vectors.
#[derive(Debug, Clone, PartialEq, Default)]
pub struct DVector(pub Vec<Real>);

impl Deref for DVector {
    type Target = Vec<Real>;

    fn deref(&self) -> &Vec<Real> {
        &self.0
    }
}

impl DerefMut for DVector {
    fn deref_mut(&mut self) -> &mut Vec<Real> {
        &mut self.0
    }
}

impl fmt::Display for DVector {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "(")?;
        for (i, x) in self.iter().enumerate() {
            if i > 0 {
                write!(f, ", ")?;
            }
            write!(f, "{}", x)?
        }
        write!(f, ")")?;
        Ok(())
    }
}

impl FromIterator<Real> for DVector {
    fn from_iter<I: IntoIterator<Item = Real>>(iter: I) -> Self {
        DVector(Vec::from_iter(iter))
    }
}

impl IntoIterator for DVector {
    type Item = Real;
    type IntoIter = IntoIter<Real>;

    fn into_iter(self) -> IntoIter<Real> {
        self.0.into_iter()
    }
}

/// Type of dense or vectors.
#[derive(Debug, Clone)]
pub enum Vector {
    /// A vector with dense storage.
    Dense(DVector),

    /**
     * A vector with sparse storage.
     *
     * For each non-zero element this vector stores an index and the
     * value of the element in addition to the size of the vector.
     */
    Sparse { size: usize, elems: Vec<(usize, Real)> },
}

impl fmt::Display for Vector {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Vector::Dense(ref v) => write!(f, "{}", v),
            Vector::Sparse { size, ref elems } => {
                let mut it = elems.iter();
                write!(f, "{}:(", size)?;
                if let Some(&(i, x)) = it.next() {
                    write!(f, "{}:{}", i, x)?;
                    for &(i, x) in it {
                        write!(f, ", {}:{}", i, x)?;
                    }
                }
                write!(f, ")")
            }
        }
    }
}

impl DVector {
    /// Set all elements to 0.
    pub fn init0(&mut self, size: usize) {
        self.clear();
        self.extend((0..size).map(|_| 0.0));
        // let n = self.len();
        // self.resize(size, 0.0);
        // for i in 0..min(n, size) {
        //     unsafe { *self.get_unchecked_mut(i) = 0.0; }
        // }
    }

    /// Set self = factor * y.
    pub fn scal(&mut self, factor: Real, y: &DVector) {
        self.clear();
        self.extend(y.iter().map(|y| factor * y));
        // self.resize(y.len(), 0.0);
        // for (i, x) in y.iter().enumerate() {
        //     unsafe { *self.get_unchecked_mut(i) = factor * x; }
        // }
    }

    /// Return factor * self.
    pub fn scaled(&self, factor: Real) -> DVector {
        let mut x = Vec::with_capacity(self.len());
        x.extend(self.iter().map(|&x| x * factor));
        DVector(x)
    }

    /// Return the inner product with another vector.
    pub fn dot(&self, other: &DVector) -> Real {
        assert_eq!(self.len(), other.len());
        self.dot_begin(other)
        // let mut ip = 0.0;
        // for i in 0..self.len() {
        //     ip += unsafe { self.get_unchecked(i) * other.get_unchecked(i) };
        // }
        // return ip;
    }

    /// Return the inner product with another vector.
    ///
    /// The inner product is computed on the smaller of the two
    /// dimensions. All other elements are assumed to be zero.
    pub fn dot_begin(&self, other: &DVector) -> Real {
        self.iter().zip(other.iter()).map(|(a, b)| a * b).sum()
        // let mut ip = 0.0;
        // for i in 0..min(self.len(), other.len()) {
        //     ip += unsafe { self.get_unchecked(i) * other.get_unchecked(i) };
        // }
        // return ip;
    }

    /// Add two vectors and store result in this vector.
    pub fn add(&mut self, x: &DVector, y: &DVector) {
        assert_eq!(x.len(), y.len());
        self.clear();
        self.extend(x.iter().zip(y.iter()).map(|(a, b)| a + b));
        // self.resize(x.len(), 0.0);
        // for i in 0..x.len() {
        //     unsafe { *self.get_unchecked_mut(i) = *x.get_unchecked(i) + *y.get_unchecked(i) };
        // }
    }

    /// Add two vectors and store result in this vector.
    pub fn add_scaled(&mut self, alpha: Real, y: &DVector) {
        assert_eq!(self.len(), y.len());
        for (x, y) in self.iter_mut().zip(y.iter()) {
            *x += alpha * y;
        }
        // for i in 0..self.len() {
        //     unsafe { *self.get_unchecked_mut(i) += alpha * *y.get_unchecked(i) };
        // }
    }

    /// Add two vectors and store result in this vector.
    ///
    /// In contrast to `add_scaled`, the two vectors might have
    /// different sizes. The size of the resulting vector is the
    /// larger of the two vector sizes and the remaining entries of
    /// the smaller vector are assumed to be 0.0.
    pub fn add_scaled_begin(&mut self, alpha: Real, y: &DVector) {
        for (x, y) in self.iter_mut().zip(y.iter()) {
            *x += alpha * y;
        }
        let n = self.len();
        if n < y.len() {
            self.extend_from_slice(&y[n..]);
        }
        // if self.len() < y.len() {
        //     self.resize(y.len(), 0.0);
        // }
        // for i in 0..y.len() {
        //     unsafe { *self.get_unchecked_mut(i) += alpha * *y.get_unchecked(i) };
        // }
    }

    /// Combines this vector with another vector.
    pub fn combine(&self, self_factor: Real, other_factor: Real, other: &DVector) -> DVector {
        assert_eq!(self.len(), other.len());
        let mut result = vec![];
        result.extend(
            self.iter()
                .zip(other.iter())
                .map(|(a, b)| self_factor * a + other_factor * b),
        );
        DVector(result)
        // let mut result = DVector(Vec::with_capacity(self.len()));
        // for i in 0..self.len() {
        //     result.push(unsafe {
        //         self_factor * *self.get_unchecked(i) +
        //             other_factor * *other.get_unchecked(i)
        //     });
        // }
        // result
    }

    /// Return the 2-norm of this vector.
    pub fn norm2(&self) -> Real {
        self.iter().map(|x| x * x).sum::<Real>().sqrt()
        // let mut norm = 0.0;
        // for x in self.iter() {
        //     norm += x * x
        // }
        // norm.sqrt()
    }
}

impl Vector {
    /**
     * Return a sparse vector with the given non-zeros.
     */
    pub fn new_sparse(n: usize, indices: &[usize], values: &[Real]) -> Vector {
        assert_eq!(indices.len(), values.len());

        if indices.is_empty() {
            Vector::Sparse { size: n, elems: vec![] }
        } else {
            let mut ordered: Vec<_> = (0..n).collect();
            ordered.sort_by_key(|&i| indices[i]);
            assert!(*indices.last().unwrap() < n);
            let mut elems = Vec::with_capacity(indices.len());
            let mut last_idx = n;
            for i in ordered {
                let val = unsafe { *values.get_unchecked(i) };
                if val != 0.0 {
                    let idx = unsafe { *indices.get_unchecked(i) };
                    if idx != last_idx {
                        elems.push((idx, val));
                        last_idx = idx;
                    } else {
                        elems.last_mut().unwrap().1 += val;
                        if elems.last_mut().unwrap().1 == 0.0 {
                            elems.pop();
                            last_idx = n;
                        }
                    }
                }
            }
            Vector::Sparse { size: n, elems }
        }
    }

    /**
     * Convert vector to a dense vector.
     *
     * This function always returns a copy of the vector.
     */
    pub fn to_dense(&self) -> DVector {
        match *self {
            Vector::Dense(ref x) => x.clone(),
            Vector::Sparse { size: n, elems: ref xs } => {
                let mut v = vec![0.0; n];
                for &(i, x) in xs {
                    unsafe { *v.get_unchecked_mut(i) = x };
                }
                DVector(v)
            }
        }
    }
}
