/*
 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"
#include "femfenics_factory.h"
#include "dolfin_compat.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)
                    {
                      femfenics_factory factory;

                      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<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"); }
                        }

                      retval(0) = factory.matrix (A);
                      retval(1) = factory.vector (B);
                    }
                }
              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)
                    {
                      femfenics_factory factory;

                      dolfin::Matrix A;
                      dolfin::assemble (A, a);
                      dolfin::Vector B;
                      dolfin::assemble (B, b);
#ifdef LATEST_DOLFIN
                      dolfin::Vector x (MPI_COMM_WORLD, myx.length ());
#else
                      dolfin::Vector x (myx.length ());
#endif

                      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<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"); }
                        }

                      retval(0) = factory.matrix (A);
                      retval(1) = factory.vector (B);
                      retval(2) = factory.vector (x);
                    }
                }
              else
                { error ("assemble_system: unknown size"); }
            }
        }
    }
  return retval;
}