/*
 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_system, args, nargout,
"-*- texinfo -*-\n\
@deftypefn {Function File} {[@var{A}], [@var{b}], [@var{x}(Optional)]} = \
assemble_system (@var{form_a}, @var{form_L}, @var{DirichletBC})\n\
Construct the discretization of a system and apply essential BC.\n\
The input arguments are\n\
@itemize @bullet\n\
@item @var{form_a} which is the BilinearForm to assemble.\n\
@item @var{form_L} which is the LinearForm to assemble.\n\
@item @var{DirichletBC} represents the optional BC applied to \
the system. \n\
@end itemize \n\
The output @var{A} is a matrix representing the @var{form_a} while \
@var{b} represents @var{form_L}. \n\
If boundary conditions have to be applied to a vector for a nonlinear problem \
then it should be provide as 3rd argument and it will be given back \
as the 3rd 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 (! form_type_loaded)
    {
      form::register_type ();
      form_type_loaded = true;
      mlock ();
    }

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

  if (nargout == 2)
    {
      if (nargin < 2)
        print_usage ();
      else
        {

          if (args(0).type_id () == form::static_type_id ()
              && args(1).type_id () == form::static_type_id ())
            {
              const form & frm1 =
                static_cast<const form&> (args(0).get_rep ());
              const form & frm2 =
                static_cast<const form&> (args(1).get_rep ());

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

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

                      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, B);
                            }
                          else
                            error ("assemble_system: 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;

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

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

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

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

                  if (a.rank () == 2 && b.rank () == 1)
                    {
                      dolfin::parameters["linear_algebra_backend"] = "uBLAS";
                      dolfin::Matrix A;
                      dolfin::assemble (A, a);
                      dolfin::Vector B;
                      dolfin::assemble (B, b);
                      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 = 3; 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, B, x);

                            }
                          else
                            error ("assemble_system: 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;

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

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

                      retval(1) = myb;
                      retval(2) = myc;
                    }
                }
              else
                error ("assemble_system: unknown size");
            }
        }
    }
  return retval;
}