/*
 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}(Optional), \
                 @var{...}) \n\
The input arguments are\n\
@itemize @bullet\n\
@item @var{form a} the bilinear form to assemble.\n\
@item @var{form a} the linear form to assemble.\n\
@item @var{DirichletBC} represents the optional BC that you wish to apply to\n\
the system. If more than one BC has to be applied, just list them.\n\
@end itemize \n\
The output @var{A} is a discretized representation of the system:\n\
@itemize @bullet\n\
@item @var{A} is the sparse Matrix corresponding to the @var{form a}\n\
@item @var{A} is the Vector corresponding to the @var{form L}\n\
@end itemize \n\
If you need to apply boundary condition to a system for a nonlinear problem \n\
then you should provide as 3rd argument the vector and you will receive it \n\
back as the third output argument.\n\
For an example of this situation, you can look the example HyperElasticity.m\n\
@seealso{BilinearForm, LinearForm, ResidualForm, JacobianForm}\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;
}