/*
 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 "form.h"
#include "boundarycondition.h"

DEFUN_DLD (assemble, args, nargout, 
"-*- texinfo -*-\n\
@deftypefn {Function File} {[@var{A}], [@var{x}(Optional)]} = \
assemble (@var{form_a}, @var{DirichletBC}) \n\
Construct the discretization of a Form and apply essential BC.\n\
The input arguments are\n\
@itemize @bullet\n\
@item @var{form_a} which is the form to assemble.\n\
It can be a form of rank 2 (BilinearForm or JacobianForm), \
a form of rank 1 (LinearForm or ResidualForm) or a form \
of rank 0 (Functional).\n\
@item @var{DirichletBC} represents the optional BC applied to \
the system. \n\
@end itemize \n\
The output @var{A} is a discretized representation of the @var{form_a}:\n\
@itemize @bullet\n\
@item @var{A} is a sparse Matrix if @var{form_a} is a bilinear form\n\
@item @var{A} is a Vector if @var{form_a} is a linear form\n\
@item @var{A} is a Double if @var{form_a} is a functional\n\
@end itemize \n\
If boundary condition has to be applied to a vector for a nonlinear problem \
then it should be provide as 2nd argument and it will be given back \
as the second output argument. For an example of this situation, please refer \
to the HyperElasticity example. \n\
@seealso{BilinearForm, LinearForm, ResidualForm, JacobianForm, Functional}\n\
@end deftypefn")
{
  int nargin = args.length ();
  octave_value_list retval;

  if (! boundarycondition_type_loaded)
    {
      boundarycondition::register_type ();
      boundarycondition_type_loaded = true;
      mlock ();
    }

  if (! form_type_loaded)
    {
      form::register_type ();
      form_type_loaded = true;
      mlock ();
    }

  if (nargout == 1)
    {
      if (nargin < 1)
        print_usage ();
      else
        {
          if (args(0).type_id () == form::static_type_id ())
            {
              const form & frm = static_cast<const form&> (args(0).get_rep ());

              if (! error_state)
                {
                  const dolfin::Form & a = frm.get_form ();
                  a.check ();

                  if (a.rank () == 2)
                    {
                      dolfin::parameters["linear_algebra_backend"] = "uBLAS";
                      dolfin::Matrix A;
                      dolfin::assemble (A, a);

                      for (std::size_t i = 1; i < nargin; ++i)
                        {
                          if (args(i).type_id () == 
                              boundarycondition::static_type_id ())
                            {
                              const boundarycondition & bc
                                = static_cast<const boundarycondition&> 
                                (args(i).get_rep ());

                              const 
                                std::vector<boost::shared_ptr
                                            <const dolfin::DirichletBC> > 
                                & pbc = bc.get_bc ();
                              
                              for (std::size_t j = 0; j < pbc.size (); ++j)
                                pbc[j]->apply(A);
                            }
                          else
                            error ("assemble: unknown argument type");
                        }

                      // Get capacity of the dolfin sparse matrix
                      boost::tuples::tuple<const std::size_t*, 
                                           const std::size_t*, 
                                           const double*, int> 
                        aa = A.data ();

                      int nnz = aa.get<3> ();
                      std::size_t nr = A.size (0), nc = A.size (1);
                      std::vector<double> data_tmp;
                      std::vector<std::size_t> cidx_tmp;

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

                      octave_idx_type* orow = ridx.fortran_vec ();
                      octave_idx_type* oc = cidx.fortran_vec ();
                      double* ov = data.fortran_vec ();

                      for (std::size_t i = 0; i < nr; ++i)
                       {
                         A.getrow (i, cidx_tmp, data_tmp);
                         nz += cidx_tmp.size ();

                         for (octave_idx_type j = 0; 
                              j < cidx_tmp.size (); ++j)
                           {
                             orow [ii + j] = i;
                             oc [ii + j] = cidx_tmp [j];
                             ov [ii + j] = data_tmp [j];
                           }

                         ii = nz;
                       }

                      dims(0) = ii;
                      ridx.resize (dims);
                      cidx.resize (dims);
                      data.resize (dims);

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

                  else if (a.rank () == 1)
                    {
                      dolfin::Vector A;
                      dolfin::assemble (A, a);

                      for (std::size_t i = 1; i < nargin; ++i)
                        {
                          if (args(i).type_id () == 
                              boundarycondition::static_type_id ())
                            {
                              const boundarycondition & bc
                                = static_cast<const boundarycondition&> 
                                (args(i).get_rep ());

                              const std::vector<boost::shared_ptr 
                                                <const dolfin::DirichletBC> >
                                & pbc = bc.get_bc ();

                              for (std::size_t j = 0; j < pbc.size (); ++j)
                                pbc[j]->apply(A);
                            }
                          else
                            error ("assemble: unknown argument type");
                        }

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

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

                      retval(0) = myb;
                    }

                  else if (a.rank () == 0)
                    {
                      double b = dolfin::assemble (a);
                      retval(0) = octave_value (b);
                    }

                  else
                    error ("assemble: unknown form size");
                }
            }
        }
    }
  else if (nargout == 2)
    {
      if (nargin < 2)
        print_usage ();
      else
        {
          if (args(0).type_id () == form::static_type_id ())
            {
              const form & frm = 
                static_cast<const form &> (args(0).get_rep ());
              const Array<double> myx = args(1).array_value ();

              if (! error_state)
                {
                  const dolfin::Form & a = frm.get_form ();
                  a.check ();

                  if (a.rank () == 1)
                    {
                      dolfin::Vector A;
                      dolfin::assemble (A, a);

                      dolfin::Vector x (myx.length ());
                      for (std::size_t i = 0; i < myx.length (); ++i)
                        x.setitem (i, myx.xelem (i));

                      for (std::size_t i = 2; i < nargin; ++i)
                        {
                          if (args(i).type_id () == 
                              boundarycondition::static_type_id ())
                            {
                              const boundarycondition & bc
                                = static_cast<const boundarycondition&> 
                                (args(i).get_rep ());

                              const std::vector<boost::shared_ptr 
                                                <const dolfin::DirichletBC> > 
                                & pbc = bc.get_bc ();

                              for (std::size_t j = 0; 
                                   j < pbc.size (); ++j)
                                pbc[j]->apply(A, x);
                            }
                          else
                            error ("assemble: unknown argument type");
                        }

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

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

                      retval(0) = myb;
                      retval(1) = myc;
                    }

                  else
                    error ("assemble: unknown size");
                }
            }
        }
    }

  return retval;
}