/*
 Copyright (C) 2014 Eugenio Gianniti <eugenio.gianniti@mail.polimi.it>

 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/>.
*/

#include "PETSc_factory.h"
#ifdef LATEST_DOLFIN
#include <boost/mpi.hpp>
#include <boost/serialization/vector.hpp>
#endif

femfenics_base_factory const &
PETSc_factory::instance (void)
{
  static PETSc_factory const theinstance;
  return theinstance;
}

octave_value
PETSc_factory::matrix (dolfin::Matrix const & A) const
{
  octave_value retval;

  unsigned const size =
#ifdef LATEST_DOLFIN
    dolfin::MPI::size (MPI_COMM_WORLD);
#else
    dolfin::MPI::num_processes ();
#endif

  if (size > 1)
    { retval = do_matrix_parallel (A); }
  else
    { retval = do_matrix_serial (A); }

  return retval;
}

octave_value
PETSc_factory::vector (dolfin::Vector const & b) const
{
  octave_value retval;

  unsigned const size =
#ifdef LATEST_DOLFIN
    dolfin::MPI::size (MPI_COMM_WORLD);
#else
    dolfin::MPI::num_processes ();
#endif

  if (size > 1)
    { retval = do_vector_parallel (b); }
  else
    { retval = do_vector_serial (b); }

  return retval;
}

void
PETSc_factory::add_to_arrays (Array <octave_idx_type> & ridx,
                              Array <octave_idx_type> & cidx,
                              Array <double> & values,
                              std::vector <std::size_t> const & row,
                              std::vector <std::size_t> const & col,
                              std::vector <double> const & val)
{
  octave_idx_type pos = cidx.numel ();

  dim_vector dims = cidx.dims ();
  if (dims(1) == 1)
    { dims(0) += col.size (); }
  else
    { dims(1) += col.size (); }
  ridx.resize (dims, 0);
  cidx.resize (dims, 0);
  values.resize (dims, 0);

  std::vector <std::size_t>::const_iterator rit = row.begin (),
                                            cit = col.begin ();
  std::vector <double>::const_iterator vit = val.begin ();

  while (cit != col.end ())
    {
      ridx(pos) = *rit;
      cidx(pos) = *cit;
      values(pos) = *vit;

      ++rit;
      ++cit;
      ++vit;
      ++pos;
    }

  return;
}

octave_value
PETSc_factory::do_matrix_serial (dolfin::Matrix const & A) const
{
  octave_value retval;

  std::size_t nr = A.size (0), nc = A.size (1);

  dim_vector dims (0, 1);
  Array <octave_idx_type> ridx (dims, 0), cidx (dims, 0);
  Array<double> data (dims, 0);

  for (std::size_t i = 0; i < nr; ++i)
    {
      std::vector <double> data_tmp;
      std::vector <std::size_t> cidx_tmp;
      A.getrow (i, cidx_tmp, data_tmp);
      std::vector <std::size_t> ridx_tmp (cidx_tmp.size (), i);
      add_to_arrays (ridx, cidx, data, ridx_tmp, cidx_tmp, data_tmp);
    }

  SparseMatrix sm (data, ridx, cidx, nr, nc);
  retval = sm;

  return retval;
}

octave_value
PETSc_factory::do_matrix_parallel (dolfin::Matrix const & A) const
{
  octave_value retval;

  std::vector <std::size_t> rows, cols;
  std::vector <double> vals;
  std::size_t const nrows = A.size (0), ncols = A.size (1);

  std::pair <std::size_t, std::size_t> const range = A.local_range (0);
  for (std::size_t i = range.first; i < range.second; ++i)
    {
      std::vector <std::size_t> cols_tmp;
      std::vector <double> vals_tmp;
      A.getrow (i, cols_tmp, vals_tmp);
      cols.insert (cols.end (), cols_tmp.begin (), cols_tmp.end ());
      vals.insert (vals.end (), vals_tmp.begin (), vals_tmp.end ());
      for (std::size_t j = 0; j < cols_tmp.size (); ++j)
        { rows.push_back (i); }
    }

  boost::mpi::communicator world;

  unsigned const size = world.size ();
  unsigned const rank = world.rank ();

  for (unsigned p = 1; p < size; ++p)
    {
      if (rank == p)
        {
          unsigned const tag = 3 * (p - 1);
          world.send (0, tag, rows);
          world.send (0, tag + 1, cols);
          world.send (0, tag + 2, vals);
        }
    }

  if (rank == 0)
    {
      dim_vector dims (0, 1);
      Array <octave_idx_type> ridx (dims, 0), cidx (dims, 0);
      Array <double> data (dims, 0);
      add_to_arrays (ridx, cidx, data, rows, cols, vals);

      for (unsigned p = 1; p < size; ++p)
        {
          std::vector <double> new_vals;
          std::vector <std::size_t> new_rows, new_cols;
          unsigned const tag = 3 * (p - 1);

          world.recv (p, tag, new_rows);
          world.recv (p, tag + 1, new_cols);
          world.recv (p, tag + 2, new_vals);

          add_to_arrays (ridx, cidx, data, new_rows, new_cols, new_vals);
        }

      SparseMatrix sm (data, ridx, cidx, nrows, ncols);
      retval = sm;
    }
  else
    // Return an identity matrix just to avoid warnings or errors
    {
      dim_vector dims (nrows, 1);
      Array <octave_idx_type> ridx (dims), cidx (dims);
      Array <double> data (dims, 1);

      for (std::size_t i = 0; i < nrows; ++i)
        {
          ridx.xelem (i) = i;
          cidx.xelem (i) = i;
        }

      SparseMatrix sm (data, ridx, cidx, nrows, ncols);
      retval = sm;
    }

  return retval;
}

octave_value
PETSc_factory::do_vector_serial (dolfin::Vector const & b) const
{
  octave_value retval;

  dim_vector dims;
  dims.resize (2);
  dims(0) = b.size ();
  dims(1) = 1;
  Array <double> myb (dims);

  for (std::size_t i = 0; i < b.size (); ++i)
    { myb.xelem (i) = b[i]; }

  retval = myb;

  return retval;
}

octave_value
PETSc_factory::do_vector_parallel (dolfin::Vector const & b) const
{
  octave_value retval;

  std::size_t const length = b.size ();
  dim_vector dims (length, 1);
  Array <double> vec (dims, 0);

  std::vector <double> entries;
  b.gather_on_zero (entries);

#ifdef LATEST_DOLFIN
  if (dolfin::MPI::rank (MPI_COMM_WORLD) == 0)
#else
  if (dolfin::MPI::process_number () == 0)
#endif
    for (std::size_t i = 0; i < length; ++i)
      { vec.xelem (i) = entries[i]; }

  // On nodes other than zero, return an all zero vector of the right length
  retval = vec;

  return retval;
}