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

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

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

  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);
                  // nz shoud be estimated in a better way
                  double alpha = 0.0005;
                  octave_idx_type nz = alpha * nr * nc;
                  SparseMatrix sm (nr, nc, nz);

                  std::size_t ii = 0;
                  sm.cidx (0) = 0;
                  for (std::size_t j = 0; j < nc; ++j)
                   {
                     for (std::size_t i = 0; i < nr; ++i)
                       {
                         double tmp = A(i, j);
                         if (tmp != 0.)
                           {
                             if (ii == nz)
                               {
                                 nz = 1.2 * ((nc * ii) / (j + 1));
                                 sm.change_capacity (nz);
                               }
                             sm.data(ii) = tmp;
                             sm.ridx(ii) = i;
                             ++ii;
                           }
                       }
                     sm.cidx(j+1) = ii;
                  }
                  sm.maybe_compress ();
                  retval = 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(i) = A[i];

                  retval = myb;
                }

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