///
/// This file is part of Rheolef.
///
/// Copyright (C) 2000-2009 Pierre Saramito <Pierre.Saramito@imag.fr>
///
/// Rheolef 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 2 of the License, or
/// (at your option) any later version.
///
/// Rheolef 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 Rheolef; if not, write to the Free Software
/// Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
/// 
/// =========================================================================

#include "rheolef/csr.h"
#include "rheolef/asr.h"
#include "rheolef/asr_to_csr.h"
#include "rheolef/msg_util.h"
#include "rheolef/csr_amux.h"
#include "rheolef/csr_cumul_trans_mult.h"
using namespace std;
namespace rheolef {
// ----------------------------------------------------------------------------
// class member functions
// ----------------------------------------------------------------------------

template<class T>
csr_rep<T,sequential>::csr_rep(const distributor& row_ownership, const distributor& col_ownership, size_type nnz1) 
 : vector_of_iterator<pair_type>(row_ownership.size()+1),
   _row_ownership (row_ownership),
   _col_ownership (col_ownership),
   _data (nnz1),
   _is_symmetric (false),
   _is_definite_positive (false),
   _pattern_dimension (0)
{
}
template<class T>
void
csr_rep<T,sequential>::resize (const distributor& row_ownership, const distributor& col_ownership, size_type nnz1) 
{
   vector_of_iterator<pair_type>::resize (row_ownership.size()+1);
   _row_ownership = row_ownership;
   _col_ownership = col_ownership;
   _data.resize (nnz1);
   // first pointer points to the beginning of the data:
   vector_of_iterator<pair_type>::operator[](0) = _data.begin().operator->();
}
template<class T>
csr_rep<T,sequential>::csr_rep(size_type loc_nrow1, size_type loc_ncol1, size_type loc_nnz1) 
 : vector_of_iterator<pair_type> (loc_nrow1+1),
   _row_ownership (distributor::decide, communicator(), loc_nrow1),
   _col_ownership (distributor::decide, communicator(), loc_ncol1),
   _data (loc_nnz1),
   _is_symmetric (false),
   _is_definite_positive (false),
   _pattern_dimension (0)
{
}
template<class T>
void
csr_rep<T,sequential>::resize (size_type loc_nrow1, size_type loc_ncol1, size_type loc_nnz1)
{
   vector_of_iterator<pair_type>::resize (loc_nrow1+1);
   _row_ownership = distributor (distributor::decide, communicator(), loc_nrow1);
   _col_ownership = distributor (distributor::decide, communicator(), loc_ncol1);
   _data.resize (loc_nnz1);
   // first pointer points to the beginning of the data:
   vector_of_iterator<pair_type>::operator[](0) = _data.begin().operator->();
}
template<class T>
csr_rep<T,sequential>::csr_rep(const csr_rep<T,sequential>& b)
 : vector_of_iterator<pair_type>(b.nrow()+1),
   _row_ownership (b.row_ownership()),
   _col_ownership (b.col_ownership()),
   _data(b._data),
   _is_symmetric (b._is_symmetric),
   _is_definite_positive (b._is_definite_positive),
   _pattern_dimension (b._pattern_dimension)
{
  // physical copy of csr
  const_iterator ib = b.begin();
  iterator       ia = begin();
  ia[0] = _data.begin().operator->();
  for (size_type i = 0, n = b.nrow(); i < n; i++) {
    ia [i+1] = ia[0] + (ib[i+1] - ib[0]);
  }
}
template<class T>
template<class A>
void
csr_rep<T,sequential>::build_from_asr (const asr<T,sequential,A>& a)
{
  resize (a.row_ownership(), a.col_ownership(), a.nnz());
  typedef std::pair<size_type,T>                pair_type;
  typedef std::pair<size_type,T>                const_pair_type;
  //typedef typename asr<T>::row_type::value_type const_pair_type;
  asr_to_csr (
	a.begin(),
        a.end(), 
	always_true<const_pair_type>(), 
	pair_identity<const_pair_type,pair_type>(), 
        vector_of_iterator<pair_type>::begin(), 
        _data.begin());
} 
template<class T>
odiststream&
csr_rep<T,sequential>::put_matrix_market (odiststream& ods, size_type first_dis_i) const
{
  std::ostream& os = ods.os();
  os << setprecision (std::numeric_limits<T>::digits10)
     << "%%MatrixMarket matrix coordinate real general" << std::endl
     << nrow() << " " << ncol() << " " << nnz() << endl;
  const_iterator ia = begin();
  const size_type base = 1;
  for (size_type i = 0, n = nrow(); i < n; i++) {
    for (const_data_iterator iter = ia[i], last = ia[i+1]; iter != last; ++iter) {
      os << i+first_dis_i+base << " "
         << (*iter).first+base << " "
         << (*iter).second << endl;	
    }
  }
  return ods;
}
template<class T>
odiststream&
csr_rep<T,sequential>::put_sparse_matlab (odiststream& ods, size_type first_dis_i) const
{
  std::ostream& os = ods.os();
  std::string name = iorheo::getbasename(ods.os());
  if (name == "") name = "a";
  os << setprecision (std::numeric_limits<T>::digits10)
     << name << " = spalloc(" << nrow() << "," << ncol() << "," << nnz() << ");" << endl;
  const_iterator ia = begin();
  const size_type base = 1;
  for (size_type i = 0, n = nrow(); i < n; i++) {
    for (const_data_iterator iter = ia[i], last = ia[i+1]; iter != last; ++iter) {
      os << name << "(" << i+first_dis_i+base << "," << (*iter).first+base << ") = "
         << (*iter).second << ";" << endl;	
    }
  }
  return ods;
}
template<class T>
odiststream&
csr_rep<T,sequential>::put (odiststream& ods, size_type first_dis_i) const
{
    iorheo::flag_type format = iorheo::flags(ods.os()) & iorheo::format_field;
    if (format [iorheo::matlab] || format [iorheo::sparse_matlab]) 
					{ return put_sparse_matlab   (ods,first_dis_i); }
    // default is matrix market format
    return put_matrix_market (ods,first_dis_i);
}
template<class T>
void
csr_rep<T,sequential>::dump (const string& name, size_type first_dis_i) const
{
    std::ofstream os (name.c_str());
    std::cerr << "! file \"" << name << "\" created." << std::endl;
    odiststream ods(os);
    put (ods);
}
// ----------------------------------------------------------------------------
// basic linear algebra
// ----------------------------------------------------------------------------

template<class T>
void
csr_rep<T,sequential>::mult(
    const vec<T,sequential>& x,
    vec<T,sequential>&       y)
    const
{
    check_macro (x.size() == ncol(), "csr*vec: incompatible csr("<<nrow()<<","<<ncol()<<") and vec("<<x.size()<<")");
    csr_amux (
        vector_of_iterator<pair_type>::begin(), 
        vector_of_iterator<pair_type>::end(), 
        x.begin(), 
        details::generic_set_op(),
        y.begin());
}
template<class T>
void
csr_rep<T,sequential>::trans_mult(
    const vec<T,sequential>& x,
    vec<T,sequential>&       y)
    const
{
    // reset y and then cumul
    check_macro (x.size() == nrow(), "csr.trans_mult(vec): incompatible csr("<<nrow()<<","<<ncol()<<") and vec("<<x.size()<<")");
    std::fill (y.begin(), y.end(), T(0));
    csr_cumul_trans_mult (
        vector_of_iterator<pair_type>::begin(), 
        vector_of_iterator<pair_type>::end(), 
        x.begin(), 
        details::generic_set_plus_op(),
        y.begin());
}
template<class T>
csr_rep<T,sequential>&
csr_rep<T,sequential>::operator*= (const T& lambda)
{
  iterator ia = begin();
  for (data_iterator p = ia[0], last_p = ia[nrow()]; p != last_p; p++) {
    (*p).second *= lambda;
  }
  return *this;
}
// ----------------------------------------------------------------------------
// expression c=a+b and c=a-b with a temporary c=*this
// ----------------------------------------------------------------------------
template<class T>
template<class BinaryOp>
void
csr_rep<T,sequential>::assign_add (
    const csr_rep<T,sequential>& a, 
    const csr_rep<T,sequential>& b,
    BinaryOp binop)
{
    check_macro (a.nrow() == b.nrow() && a.ncol() == b.ncol(),
	"incompatible csr add(a,b): a("<<a.nrow()<<":"<<a.ncol()<<") and "
	"b("<<b.nrow()<<":"<<b.ncol()<<")");
    //
    // first pass: compute nnz_c and resize
    //
    size_type nnz_c = 0;
    const size_type infty = std::numeric_limits<size_type>::max();
    const_iterator ia = a.begin();
    const_iterator ib = b.begin();
    for (size_type i = 0, n = a.nrow(); i < n; i++) {
        for (const_data_iterator iter_jva = ia[i], last_jva = ia[i+1],
                                 iter_jvb = ib[i], last_jvb = ib[i+1];
	    iter_jva != last_jva || iter_jvb != last_jvb; ) {

            size_type ja = iter_jva == last_jva ? infty : (*iter_jva).first;
            size_type jb = iter_jvb == last_jvb ? infty : (*iter_jvb).first;
	    if (ja == jb) {
		iter_jva++;
		iter_jvb++;
	    } else if (ja < jb) {
		iter_jva++;
            } else {
		iter_jvb++;
            }
	    nnz_c++;
  	}
    }
    resize (a.row_ownership(), b.col_ownership(), nnz_c);
    data_iterator iter_jvc = _data.begin().operator->();
    iterator ic = begin();
    *ic++ = iter_jvc;
    //
    // second pass: add and store in c
    //
    for (size_type i = 0, n = a.nrow(); i < n; i++) {
        for (const_data_iterator iter_jva = ia[i], last_jva = ia[i+1],
                                 iter_jvb = ib[i], last_jvb = ib[i+1];
	    iter_jva != last_jva || iter_jvb != last_jvb; ) {

            size_type ja = iter_jva == last_jva ? infty : (*iter_jva).first;
            size_type jb = iter_jvb == last_jvb ? infty : (*iter_jvb).first;
	    if (ja == jb) {
		*iter_jvc++ = std::pair<size_type,T> (ja, binop((*iter_jva).second, (*iter_jvb).second));
		iter_jva++;
		iter_jvb++;
	    } else if (ja < jb) {
		*iter_jvc++ = std::pair<size_type,T> (ja, binop((*iter_jva).second, T(0)));
		iter_jva++;
            } else {
		*iter_jvc++ = std::pair<size_type,T> (jb, binop(T(0), (*iter_jvb).second));
		iter_jvb++;
            }
  	}
        *ic++ = iter_jvc;
    }
    set_symmetry          (a.is_symmetric() && b.is_symmetric());
    set_pattern_dimension (std::max(a.pattern_dimension(), b.pattern_dimension()));
    set_definite_positive (a.is_definite_positive() || b.is_definite_positive());
    // perhaps too optimist for definite_positive: should be ajusted before calling a solver
}
// ----------------------------------------------------------------------------
// trans(a)
// ----------------------------------------------------------------------------
/*@! 
 \vfill \pagebreak \mbox{} \vfill \begin{algorithm}[h]
  \caption{{\tt sort}: sort rows by increasing column order}
  \begin{algorithmic}
    \INPUT {the matrix in CSR format}
      ia(0:nrow), ja(0:nnz-1), a(0:nnz-1)
    \ENDINPUT
    \OUTPUT {the sorted CSR matrix}
      iao(0:nrow), jao(0:nnzl-1), ao(0:nnzl-1)
    \ENDOUTPUT
    \BEGIN 
      {\em first: reset iao} \\
      \FORTO {i := 0} {nrow}
	  iao(i) := 0;
      \ENDFOR
	
      {\em second: compute lengths from pointers} \\
      \FORTO {i := 0} {nrow-1}
        \FORTO {p := ia(i)} {ia(i+1)-1}
	    iao (ja(p)+1)++;
        \ENDFOR
      \ENDFOR

      {\em third: compute pointers from lengths} \\
      \FORTO {i := 0} {nrow-1}
	  iao(i+1) += iao(i)
      \ENDFOR

      {\em fourth: copy values} \\
      \FORTO {i := 0} {nrow-1}
        \FORTO {p := ia(i)} {ia(i+1)-1}
          j := ja(p) \\
	  q := iao(j) \\
	  jao(q) := i \\
	  ao(q) := a(p) \\
	  iao (j)++
        \ENDFOR
      \ENDFOR

      {\em fiveth: reshift pointers} \\
      \FORTOSTEP {i := nrow-1} {0} {-1}
	iao(i+1) := iao(i);
      \ENDFOR
      iao(0) := 0
    \END
 \end{algorithmic} \end{algorithm}
 \vfill \pagebreak \mbox{} \vfill
*/

// NOTE: transposed matrix has always rows sorted by increasing column indexes
//       even if original matrix has not
template<class T>
void
csr_rep<T,sequential>::build_transpose (csr_rep<T,sequential>& b) const
{
  b.resize (col_ownership(), row_ownership(), nnz()); 

  // first pass: set ib(*) to ib(0)
  iterator ib = b.begin();
  for (size_type j = 0, m = b.nrow(); j < m; j++) {
    ib[j+1] = ib[0];
  }
  // second pass: compute lengths of row(i) of a^T in ib(i+1)-ib(0)
  const_iterator ia = begin();
  for (size_type i = 0, n = nrow(); i < n; i++) {
    for (const_data_iterator p = ia[i], last_p = ia[i+1]; p != last_p; p++) {
      size_type j = (*p).first;
      ib [j+1]++;
    }
  }
  // third pass: compute pointers from lengths
  for (size_type j = 0, m = b.nrow(); j < m; j++) {
    ib [j+1] += (ib[j]-ib[0]);
  }
  // fourth pass: store values
  data_iterator q0 = ib[0];
  for (size_type i = 0, n = nrow(); i < n; i++) {
    for (const_data_iterator p = ia[i], last_p = ia[i+1]; p != last_p; p++) {
      size_type j = (*p).first;
      data_iterator q = ib[j];
      (*q).first  = i;
      (*q).second = (*p).second;
      ib[j]++;
    }
  }
  // fiveth: shift pointers
  for (long int j = b.nrow()-1; j >= 0; j--) {
    ib[j+1] = ib[j];
  }
  ib[0] = q0;
}
// ----------------------------------------------------------------------------
// set symmetry by check
// ----------------------------------------------------------------------------
template<class T>
void
csr_rep<T,sequential>::set_symmetry_by_check (const T& tol) const
{
  if (nrow() != ncol()) {
    set_symmetry(false);
    return;
  }
  // check if a-trans(a) == 0 up to machine prec
  csr_rep<T,sequential> at, d;
  build_transpose (at);
  d.assign_add (*this, at, std::minus<T>());
  set_symmetry(d.max_abs() <= tol);
}
// ----------------------------------------------------------------------------
// expression c=a*b with a temporary c=*this
// ----------------------------------------------------------------------------
// compute the sparse size of c=a*b
template<class T>
typename csr_rep<T,sequential>::size_type
csr_csr_mult_size (
    const csr_rep<T,sequential>& a,
    const csr_rep<T,sequential>& b)
{
  typedef typename csr_rep<T,sequential>::size_type size_type;
  typename csr_rep<T,sequential>::const_iterator ia = a.begin(), ib = b.begin();
  size_type nnz_c = 0;
  for (size_type i = 0, n = a.nrow(); i < n; i++) {
    std::set<size_type> y;
    for (typename csr<T>::const_data_iterator ap = ia[i], last_ap = ia[i+1]; ap != last_ap; ++ap) {
      size_type j = (*ap).first;
      typename std::set<size_type>::iterator iter_y = y.begin();
      for (typename csr<T>::const_data_iterator bp = ib[j], last_bp = ib[j+1]; bp != last_bp; ++bp) {
        size_type k = (*bp).first;
	iter_y = y.insert(iter_y, k);
      }
    }
    nnz_c += y.size();
  }
  return nnz_c;
}
// compute the values of the sparse matrix c=a*b
template<class T>
static
void
csr_csr_mult (
    const csr_rep<T,sequential>& a,
    const csr_rep<T,sequential>& b,
          csr_rep<T,sequential>& c)
{
  typedef typename csr_rep<T,sequential>::size_type      size_type;
  typename csr_rep<T,sequential>::const_iterator ia = a.begin(), ib = b.begin();
  typename csr_rep<T,sequential>::iterator       ic = c.begin();
  for (size_type i = 0, n = a.nrow(); i < n; i++) {
    std::map<size_type,T> y;
    for (typename csr<T>::const_data_iterator ap = ia[i], last_ap = ia[i+1]; ap != last_ap; ++ap) {
      size_type j    = (*ap).first;
      const T&  a_ij = (*ap).second;
      typename std::map<size_type,T>::iterator prev_y = y.begin();
      for (typename csr<T>::const_data_iterator bp = ib[j], last_bp = ib[j+1]; bp != last_bp; ++bp) {
        size_type k    = (*bp).first;
        const T&  b_jk = (*bp).second;
	T value = a_ij*b_jk;
        typename std::map<size_type,T>::iterator curr_y = y.find(k);
        if (curr_y == y.end()) { 
          // create a new (i,k,value) entry in matrix C
	  y.insert (prev_y, std::pair<size_type,T>(k, value));
        } else {
          // increment an existing (i,k,value) entry in matrix C
          (*curr_y).second += value;
          prev_y = curr_y;
        }
      }
    }
    ic[i+1] = ic[i] + y.size();
    // copy the y temporary into the i-th row of C:
    typename std::map<size_type,T>::const_iterator iter_y = y.begin();
    for (typename csr<T>::data_iterator cp = ic[i], last_cp = ic[i+1]; cp != last_cp; ++cp, ++iter_y) {
	*cp = *iter_y;
    }
  }
  // propagate flags ; cannot stat yet for symmetry and positive_definite
  c.set_pattern_dimension (std::max(a.pattern_dimension(), b.pattern_dimension()));
}
template<class T>
void
csr_rep<T,sequential>::assign_mult (
    const csr_rep<T,sequential>& a, 
    const csr_rep<T,sequential>& b)
{
  size_type nnz_c = csr_csr_mult_size (a, b);
  resize (a.row_ownership(), b.col_ownership(), nnz_c);
  csr_csr_mult (a, b, *this);
}
// ----------------------------------------------------------------------------
// instanciation in library
// ----------------------------------------------------------------------------
#ifndef _RHEOLEF_USE_HEAP_ALLOCATOR
#define _RHEOLEF_instanciate_class(T) \
template void csr_rep<T,sequential>::build_from_asr (const asr<T,sequential,std::allocator<T> >&);
#else // _RHEOLEF_USE_HEAP_ALLOCATOR
#define _RHEOLEF_instanciate_class(T) \
template void csr_rep<T,sequential>::build_from_asr (const asr<T,sequential,std::allocator<T> >&);	\
template void csr_rep<T,sequential>::build_from_asr (const asr<T,sequential,heap_allocator<T> >&);
#endif // _RHEOLEF_USE_HEAP_ALLOCATOR

#define _RHEOLEF_instanciate(T) \
template class csr_rep<T,sequential>;					\
template void csr_rep<T,sequential>::assign_add (			\
	const csr_rep<T,sequential>&,					\
	const csr_rep<T,sequential>&,					\
	      std::plus<T>);						\
template void csr_rep<T,sequential>::assign_add (			\
	const csr_rep<T,sequential>&,					\
	const csr_rep<T,sequential>&,					\
	std::minus<T>);							\
_RHEOLEF_instanciate_class(T)

_RHEOLEF_instanciate(Float)

} // namespace rheolef