'''
 Copyright (C) 2013 - Juan Pablo Carbajal
 Copyright (C) 2013 - Marco Vassallo

 This progrm 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/>.
'''

# Author: Juan Pablo Carbajal <ajuanpi+dev@gmail.com>
# Modified by: Marco Vassallo

from io import open
from string import Template

class Makefile ():

  def __init__(self,**kwargs):
    """ Creates a makefile with the names of the files to be compiled.

    Inuts must be all strings.
    """
    if map(type,kwargs.values()) != [str]*len(kwargs):
      raise ValueError("Makefile accepts only str arguments.")

    self.name = "Makefile_"+kwargs["name"]
    self.body = Template ("""
DIR = $path
CPPFLAGS= -Wall
LIBS= -ldolfin
MKOCTFILE = mkoctfile
FFC = ffc

OCTFILES = fem_fs_$name.oct

all : $(OCTFILES) 
fs : fem_fs_$name.oct
rhs : fem_rhs_$name.oct
lhs : fem_lhs_$name.oct

fem_fs_$name.o: fem_fs_$name.cc
	$(MKOCTFILE) -I$(DIR) -I. $(CPPFLAGS) $< -c $@

fem_fs_$name.oct: $name.h fem_fs_$name.o
	$(MKOCTFILE) -s fem_fs_$name.o $(DIR)fem_init_env.o $(LIBS) -o $@

fem_rhs_$name.o: fem_rhs_$name.cc
	$(MKOCTFILE) -I$(DIR) -I. $(CPPFLAGS) $< -c $@

fem_rhs_$name.oct: $name.h fem_rhs_$name.o
	$(MKOCTFILE) -s fem_rhs_$name.o $(DIR)fem_init_env.o $(LIBS) -o $@

fem_lhs_$name.o: fem_lhs_$name.cc
	$(MKOCTFILE) -I$(DIR) -I. $(CPPFLAGS) $< -c $@

fem_lhs_$name.oct: $name.h fem_lhs_$name.o
	$(MKOCTFILE) -s fem_lhs_$name.o $(DIR)fem_init_env.o $(LIBS) -o $@

$name.h: $name.ufl
	$(FFC) -l dolfin $name.ufl

.PHONY: clean

clean:
	 rm -f fem_fs_$name.o fem_fs_$name.cc $name.h
	 rm -f fem_rhs_$name.o fem_rhs_$name.cc
	 rm -f fem_lhs_$name.o fem_lhs_$name.cc
	 rm -f $makefile

""").safe_substitute(name=kwargs["name"],
                     makefile=self.name,
                     path=kwargs["path"])

  def writeout (self):
    mfile = open (self.name, 'w')
    mfile.write(unicode(self.body))
    mfile.close()

class Functionspace ():

  def __init__(self,**kwargs):
    """ Creates a file fem_fs_$name with the name of the Variational problem to.

    Inuts must be all strings.
    """
    if map(type,kwargs.values()) != [str]*len(kwargs):
      raise ValueError("Functionspace accepts only str arguments.")

    self.name = "fem_fs_"+kwargs["name"]+".cc"
    self.body = Template ("""
#include "functionspace.h"
#include "mesh.h"
#include "$name.h"

DEFUN_DLD (fem_fs_$name, args, , "initialize a fs from a mesh declared with fem_init_mesh")
{
  int nargin = args.length ();
  octave_value retval;

  if (nargin < 1 || nargin > 1)
    print_usage ();
  else
    {

      if (! mesh_type_loaded)
        {
          mesh::register_type ();
          mesh_type_loaded = true;
          mlock ();
        }

      if (args(0).type_id () == mesh::static_type_id ())
        {
          const mesh & msho = static_cast<const mesh&> (args(0).get_rep ());
          const dolfin::Mesh & mshd = msho.get_msh ();
          boost::shared_ptr <const dolfin::FunctionSpace> g (new $name::FunctionSpace (mshd));

          if (! functionspace_type_loaded)
            {
              functionspace::register_type ();
              functionspace_type_loaded = true;
              mlock ();
            }

          retval = new functionspace(g);
        }
    }
  return retval;
}
""").safe_substitute(name=kwargs["name"])

  def writeout (self):
    mfile = open (self.name, 'w')
    mfile.write(unicode(self.body))
    mfile.close()

class RHS ():

  def __init__(self,**kwargs):
    """ Creates a file fem_rhs_$name with the name of the Variational problem to.

    Inuts must be all strings.
    """
    if map(type,kwargs.values()) != [str]*len(kwargs):
      raise ValueError("RHS accepts only str arguments.")

    self.name = "fem_rhs_"+kwargs["name"]+".cc"
    self.body = Template ("""
#include "$name.h"
#include "functionspace.h"
#include "boundarycondition.h"
#include "coefficient.h"

DEFUN_DLD (fem_rhs_$name, args, , "A = fem_rhs_$name (FUNCTIONAL SPACE, COEFF, BC)")
{

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

  if (nargin < 1)
    print_usage ();
  else
    {
      if (! functionspace_type_loaded)
        {
          functionspace::register_type ();
          functionspace_type_loaded = true;
          mlock ();
        }
      if (args(0).type_id () == functionspace::static_type_id ())
        {
          const functionspace & fspo
            = static_cast<const functionspace&> (args(0).get_rep ());

          if (! error_state)
            {
              const dolfin::FunctionSpace V = fspo.get_fsp ();
              $name::BilinearForm a (V, V);
              std::size_t ncoef = a.num_coefficients (), nc = 0;

              if (! coefficient_type_loaded)
                {
                  coefficient::register_type ();
                  coefficient_type_loaded = true;
                  mlock ();
                }

              for (std::size_t i = 1; i < nargin; ++i)
                {
                  if (args(i).type_id () == coefficient::static_type_id ())
                    {
                      const coefficient & cf
                        = static_cast <const coefficient&> (args(i).get_rep ());

                      std::size_t n = a.coefficient_number (cf.get_str ());
                      const boost::shared_ptr<const expression> & pexp = cf.get_expr ();
                      a.set_coefficient (n, pexp);
                      ++nc;
                    }
                 }

              if (nc != ncoef)
                error ("Wrong number of coefficient");
              else
                {

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

                  int nr = A.size (0), nc = A.size (1);
                  // nz shoud be estimated in a better way
                  int nz = nr * nc;
                  SparseMatrix sm (nr, nc, nz);

                  int ii = 0;
                  sm.cidx (0) = 0;
                  for (int j = 0; j < nc; ++j)
                   {
                     for (int i = 0; i < nr; i++)
                       {
                         double tmp = A(i, j);
                         if (tmp != 0.)
                           {
                             sm.data(ii) = tmp;
                             sm.ridx(ii) = i;
                             ii++;
                           }
                       }
                     sm.cidx(j+1) = ii;
                  }
                  sm.maybe_compress ();

                  retval = sm;

                }
            }
        }
    }
  return retval;
}
""").safe_substitute(name=kwargs["name"])

  def writeout (self):
    mfile = open (self.name, 'w')
    mfile.write(unicode(self.body))
    mfile.close()

class LHS ():

  def __init__(self,**kwargs):
    """ Creates a file fem_lhs_$name with the name of the Variational problem to.

    Inuts must be all strings.
    """
    if map(type,kwargs.values()) != [str]*len(kwargs):
      raise ValueError("lhs accepts only str arguments.")

    self.name = "fem_lhs_"+kwargs["name"]+".cc"
    self.body = Template ("""
#include "$name.h"
#include "functionspace.h"
#include "boundarycondition.h"
#include "coefficient.h"

DEFUN_DLD (fem_lhs_$name, args, , " b = fem_lhs_$name (FUNCTIONAL SPACE, COEFF, BC)")
{

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

  if (nargin < 1)
    print_usage ();
  else
    {

      if (! functionspace_type_loaded)
        {
          functionspace::register_type ();
          functionspace_type_loaded = true;
          mlock ();
        }
      if (args(0).type_id () == functionspace::static_type_id ())
        {
          const functionspace & fspo
            = static_cast<const functionspace&> (args(0).get_rep ());

          if (! error_state)
            {
              const dolfin::FunctionSpace V = fspo.get_fsp ();
              $name::LinearForm L (V);
              std::size_t ncoef = L.num_coefficients (), nc = 0;

              if (! coefficient_type_loaded)
                {
                  coefficient::register_type ();
                  coefficient_type_loaded = true;
                  mlock ();
                }
              for (std::size_t i = 1; i < nargin; ++i)
                {
                  if (args(i).type_id () == coefficient::static_type_id ())
                    {
                      const coefficient & cf
                        = static_cast <const coefficient&> (args(i).get_rep ());

                      std::size_t n = L.coefficient_number (cf.get_str ());
                      const boost::shared_ptr<const expression> & pexp = cf.get_expr ();
                      L.set_coefficient (n, pexp);
                      ++nc;
                    }
                 }

              if (nc != ncoef)
                error ("Wrong number of coefficient");
              else
                {

                  dolfin::Vector b;
                  dolfin::assemble (b, L);

                  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(b);
                        }
                    }

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

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

                  retval = myb;
                }
            }
        }
    }
  return retval;
}
""").safe_substitute(name=kwargs["name"])

  def writeout (self):
    mfile = open (self.name, 'w')
    mfile.write(unicode(self.body))
    mfile.close()