#include "form.h"
#include "boundarycondition.h"

DEFUN_DLD (assemble, args, nargout, "A = assemble (FORM, BC)")
{

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

  if (nargout == 1)
    {
      if (nargin < 1)
        print_usage ();
      else
        {
          if (! form_type_loaded)
            {
              form::register_type ();
              form_type_loaded = true;
              mlock ();
            }
          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::Matrix A;
                      dolfin::assemble (A, a);
                      if (! boundarycondition_type_loaded)
                        {
                          boundarycondition::register_type ();
                          boundarycondition_type_loaded = true;
                          mlock ();
                        }

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


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

                      octave_idx_type dims = A.size (0), nz = 0, ii = 0;
                      ColumnVector ridx (dims), cidx (dims), data (dims);

                      for (std::size_t i = 0; i < nr; ++i)
                       {
                         A.getrow (i, cidx_tmp, data_tmp);
                         nz += cidx_tmp.size ();
                         if (dims < nz)
                           {
                             dims = 1.2 * ((nr * nz) / (i + 1));;
                             ridx.resize (dims);
                             cidx.resize (dims);
                             data.resize (dims);
                           }

                         for (octave_idx_type j = 0; j < cidx_tmp.size (); ++j)
                           {
                             ridx.xelem (ii + j) = i + 1;
                             cidx.xelem (ii + j) = cidx_tmp [j] + 1;
                             data.xelem (ii + j) = data_tmp [j];
                           }

                         ii = nz;
                       }


                      ridx.resize (ii);
                      cidx.resize (ii);
                      data.resize (ii);

                      SparseMatrix sm (data, ridx, cidx, nr, nc);
                      retval(0) = sm;
                    }
                  else if (a.rank () == 1)
                    {
                      dolfin::Vector A;
                      dolfin::assemble (A, a);

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

                      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 size");
                }
            }
        }
    }
  else if (nargout == 2)
    {
      if (nargin < 2)
        print_usage ();
      else
        {
          if (! form_type_loaded)
            {
              form::register_type ();
              form_type_loaded = true;
              mlock ();
            }
          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));

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

                      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 NL vector: 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 NL vector: unknown size");
                }
            }
        }
    }

  return retval;
}