/*
 Copyright (C) 2013 Marco Vassallo <gedeone-octave@users.sourceforge.net>

 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 <dolfin.h>
#include "functionspace.h"
#include "function.h"
#include "dolfin_compat.h"
#ifdef LATEST_DOLFIN
#include <boost/mpi.hpp>
#include <boost/serialization/vector.hpp>
#include <boost/serialization/utility.hpp>
#endif

void copy_vector (Array <double> const &, SHARED_PTR <dolfin::GenericVector>);

DEFUN_DLD (Function, args, , "-*- texinfo -*-\n\
@deftypefn {Function File} {[@var{func}]} = \
Function (@var{name}, @var{FunctionSpace} (or @var{Function}),\
@var{Vector} (or @var{index}))\n\
Initialize an object with the values specified in a vector or extracting \
a component from a vectorial field.\n\
This function can be used in two different ways\n\
@itemize @bullet \n\
@item To create a function from a vector. In this case, the arguments are:\n\
@itemize @bullet \n\
@item @var{name} is a string representing the name of the function\n\
@item @var{FunctionSpace} is the fem-fenics function space where \
the vector is defined\n\
@item @var{Vector} specifies the values of the coefficients\
for each basisc function of the @var{FunctioSpace}\n\
@end itemize\n\
@item To extract a scalar field from a vectorial one\n\
@itemize @bullet \n\
@item @var{name} is a string representing the name of the function\n\
@item @var{Function} is the vector valued Function\n\
@item @var{Index} contains the index of the scalar field to extract.\
Index starts from 1. \n\
@end itemize\n\
@end itemize \n\
The output @var{func} is an object which contains a representation of the \
function @var{Vector} which can be plotted or saved or passed as argument \
for a variational problem. Moreover, it is possible to evaluate it on points \
of its domain.\n\
@seealso{Constant, Expression, plot, save}\n\
@end deftypefn")
{

  int nargin = args.length ();
  octave_value retval;

  if (nargin < 3 || nargin > 3)
    { print_usage (); }
  else
    {
      if (! functionspace_type_loaded)
        {
          functionspace::register_type ();
          functionspace_type_loaded = true;
          mlock ();
        }

      if (! function_type_loaded)
        {
          function::register_type ();
          function_type_loaded = true;
          mlock ();
        }

      if (args(1).type_id () == functionspace::static_type_id ())
        {
          std::string str = args(0).string_value ();
          const functionspace & fspo =
            static_cast<const functionspace &> (args(1).get_rep ());
          Array <double> myu = args(2).array_value ();

          if (!error_state)
            {
              const SHARED_PTR <const dolfin::FunctionSpace>
                & V = fspo.get_pfsp ();

              if (V->dim () != myu.length ())
                {
                  error ("The size of the functional space \
                         and of the vector must agree");
                }
              else
                {
                  dolfin::Function aux (V);
                  copy_vector (myu, aux.vector ());

                  SHARED_PTR <dolfin::Function const>
                    u (new dolfin::Function (aux));
                  retval = new function (str, u);
                }
            }
        }
      else if (args(1).type_id () == function::static_type_id ())
        {
          std::string str = args(0).string_value ();
          const function & fspo =
            static_cast<const function &> (args(1).get_rep ());
          int idx = args(2).int_value ();

          if (!error_state)
            {
              const SHARED_PTR <const dolfin::Function>
                & f = fspo.get_pfun ();

              if (f->value_rank () < 1)
                error ("Function: The function you provided\
                        isn't a vecotr field");
              else
                {
                  if (idx < 1 || idx > f->value_dimension (0))
                    { error ("Function: index out of bounds"); }
                  else
                    {
                      SHARED_PTR <const dolfin::Function>
                        u (new dolfin::Function ((*f)[idx - 1]));

                      retval = new function (str, u);
                    }
                }
            }
        }
    }
  return retval;
}

void
copy_vector (Array <double> const & arr, SHARED_PTR <dolfin::GenericVector> vec)
{
  unsigned const size =
#ifdef LATEST_DOLFIN
    dolfin::MPI::size (MPI_COMM_WORLD);
#else
    dolfin::MPI::num_processes ();
#endif

  if (size > 1)
    {
      std::vector <double> locvals;
      boost::mpi::communicator world;
      unsigned const rank = world.rank ();

      std::pair <std::size_t, std::size_t> const loc_range =
        vec->local_range ();

      for (unsigned p = 1; p < size; ++p)
        if (rank == p)
          { world.send (0, p, loc_range); }

      if (rank == 0)
        {
          for (unsigned p = 0; p < size; ++p)
            {
              std::pair <std::size_t, std::size_t> range;
              if (p == 0)
                { range = loc_range; }
              else
                { world.recv (p, p, range); }

              std::vector <double> entries;
              for (std::size_t i = range.first; i < range.second; ++i)
                { entries.push_back (arr(i)); }

              if (p == 0)
                { locvals = entries; }
              else
                { world.send (p, p, entries); }
            }
        }

      for (unsigned p = 1; p < size; ++p)
        {
          if (rank == p)
            { world.recv (0, p, locvals); }
        }

      vec->set_local (locvals);
    }
  else
    for (std::size_t i = 0; i < arr.length (); ++i)
      { vec->setitem (i, arr (i)); }

  vec->apply ("insert");

  return;
}