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

//! Finite-dimensional sparse and dense vectors.

use crate::{Aggregatable, Real};
use num_traits::Zero;
use std::borrow::Borrow;
use std::fmt;
use std::iter::FromIterator;
use std::ops::{Deref, DerefMut};
use std::vec::IntoIter;

#[cfg(feature = "blas")]
use {openblas_src as _, rs_blas as blas, std::os::raw::c_int};

/// 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));
    }

    /// Set self = factor * y.
    pub fn scal(&mut self, factor: Real, y: &DVector) {
        self.clear();
        self.extend(y.iter().map(|y| factor * y));
    }

    /// Return factor * self.
    pub fn scaled(&self, factor: Real) -> DVector {
        let mut x = DVector::default();
        x.scal(factor, self);
        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)
    }

    /// 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 {
        #[cfg(feature = "blas")]
        unsafe {
            blas::ddot(self.len().min(other.len()) as c_int, &self, 1, &other, 1)
        }
        #[cfg(not(feature = "blas"))]
        {
            self.iter().zip(other.iter()).map(|(x, y)| x * y).sum::<Real>()
        }
    }

    /// 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));
    }

    /// 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());
        #[cfg(feature = "blas")]
        unsafe {
            blas::daxpy(self.len() as c_int, alpha, &y, 1, &mut self[..], 1)
        }
        #[cfg(not(feature = "blas"))]
        {
            for (x, y) in self.iter_mut().zip(y.iter()) {
                *x += alpha * y;
            }
        }
    }

    /// 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) {
        #[cfg(feature = "blas")]
        unsafe {
            let n = self.len();
            blas::daxpy(n.min(y.len()) as c_int, alpha, &y, 1, &mut self[..], 1);
        }
        #[cfg(not(feature = "blas"))]
        {
            for (x, y) in self.iter_mut().zip(y.iter()) {
                *x += alpha * y;
            }
        }
        let n = self.len();
        if n < y.len() {
            self.extend(y[n..].iter().map(|y| alpha * y));
        }
    }

    /// Return the 2-norm of this vector.
    pub fn norm2(&self) -> Real {
        #[cfg(feature = "blas")]
        unsafe {
            blas::dnrm2(self.len() as c_int, &self, 1)
        }
        #[cfg(not(feature = "blas"))]
        {
            self.iter().map(|x| x * x).sum::<Real>().sqrt()
        }
    }
}

impl Aggregatable for DVector {
    fn new_scaled<A>(alpha: Real, other: A) -> Self
    where
        A: Borrow<Self>,
    {
        DVector::scaled(&other.borrow(), alpha)
    }

    fn add_scaled<A>(&mut self, alpha: Real, other: A)
    where
        A: Borrow<Self>,
    {
        DVector::add_scaled(self, alpha, &other.borrow())
    }
}

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 = (0..n).collect::<Vec<_>>();
            ordered.sort_unstable_by_key(|&i| indices[i]);
            assert!(indices.last().map(|&i| i < n).unwrap_or(true));
            let mut elems = Vec::with_capacity(indices.len());
            let mut last_idx = n;
            let mut last_val = Real::zero();
            for (idx, val) in ordered
                .into_iter()
                .map(|i| unsafe { (*indices.get_unchecked(i), *values.get_unchecked(i)) })
            {
                if idx != last_idx {
                    if !last_val.is_zero() {
                        elems.push((last_idx, last_val));
                    }
                    last_val = val;
                    last_idx = idx;
                } else {
                    last_val += val;
                }
            }
            if !last_val.is_zero() {
                elems.push((last_idx, last_val));
            }
            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)
            }
        }
    }
}

impl From<Vec<Real>> for DVector {
    fn from(x: Vec<Real>) -> Self {
        DVector(x)
    }
}

impl From<Vector> for DVector {
    fn from(x: Vector) -> Self {
        match x {
            Vector::Dense(x) => x,
            s => s.to_dense(),
        }
    }
}

impl From<DVector> for Vector {
    fn from(x: DVector) -> Self {
        Vector::Dense(x)
    }
}

impl From<Vec<Real>> for Vector {
    fn from(x: Vec<Real>) -> Self {
        Vector::Dense(x.into())
    }
}

#[test]
fn test_add_scaled_begin() {
    let mut x = dvec![1.0; 5];
    let y = dvec![2.0; 7];
    x.add_scaled_begin(3.0, &y);
    assert_eq!(x, dvec![7.0, 7.0, 7.0, 7.0, 7.0, 6.0, 6.0]);
}
