"""This module handles the Function class in Python.

The Function class needs special handling and is not mapped directly
by SWIG from C++ interface. Instead, a new Function class is created
which inherits both from the DOLFIN C++ Function (which is renamed to
cpp_Function in dolfin_function_pre.i) and the form compiler Function.

The resulting Function class may thus act both as a variable in a form
expression and as a Function on a mesh that may be evaluated, plotted
and stored to file.

Note that the Function class currently inherits explicitly from the FFC
Function class. In future versions, it will instead inherit from the UFL
Function class, so that DOLFIN may be form compiler agnostic."""

__author__    = "Johan Hake (hake@simula.no)"
__date__      = "2008-10-23"
__copyright__ = "Copyright (C) 2007-2008 Johan Hake"
__license__   = "GNU LGPL Version 2.1"

import types
import ffc
import dolfin
ffc_Function = ffc.Function

jit              = dolfin.jit
old_cpp_Function = dolfin.dolfin.cpp_Function

# Declarations of dummy classes
def instant(cppcode):
    class custom_cpp_Function(old_cpp_Function):
        def __init__(self,V):
            # Calling old cpp_Function
            #old_cpp_Function.__init__(self, V.Mesh(),V.DofMap(),V.compiled_form)
            self.V = V
            print "Initializing custom_cpp_Function"
        def eval(self,val,x):
            pass
        
    return custom_cpp_Function

class cpp_Function(old_cpp_Function):
    def __init__(self,V):
        # Calling old cpp_Function
        #old_cpp_Function.__init__(self, V.Mesh(),V.DofMap(),V.compiled_form)
        self._V = V

class cpp_FunctionSpace(object):
    def __init__(self, mesh, element, dofmap):
        self._mesh        = mesh
        self._ufc_element = element
        self._dofmap      = dofmap
    
    def Mesh(self):
        return self._mesh
    
    def DofMap(self):
        return self._dofmap
    
    def FiniteElement(self):
        return self._ufc_element

class FunctionSpace(cpp_FunctionSpace):
    # FIXME: Somewhat messy implementation, and probably totally wrong :P
    def __init__(self,mesh,element):
        print "Initalizing FunctionSpace"
        assert(isinstance(mesh,dolfin.Mesh))
        self._ffc_element = element
        self.mesh = mesh
        # Create simplest possible form
        if element.value_dimension(0) > 1:
            form = ffc.TestFunction(element)[0]*ffc.dx
        else:
            form = ffc.TestFunction(element)*ffc.dx
        
        self.compiled_form, module, form_data = jit(form)
        ufc_element = self.compiled_form.create_finite_element(0)
        dofmaps = dolfin.dolfin.DofMapSet(self.compiled_form, mesh)
        cpp_FunctionSpace.__init__(self, mesh, ufc_element, dofmaps.sub(0))
        
    def ffc_element(self):
        return self._ffc_element

class FunctionCreator(type):
    " Metaclass for Function"
    def __new__(cls, name, bases, dictionary):
        print "class name:",name
        print "Bases on entry:", bases
        print "Dictionary on entry:",dictionary

        # If we are creating the Function class exit without doing anything
        if name == "Function":
            return type.__new__(cls, name, bases, dictionary)
        
        # Default cpp_BaseFunction class
        cpp_BaseFunction = cpp_Function
        
        # Check the cppcode and eval attribute
        if 'cppcode' in dictionary and 'eval' in dictionary:
            raise TypeError, "Cannot create class with both 'eval' and " + \
                  "'cppcode' attributes"
        if not ('cppcode' in dictionary or 'eval' in dictionary):
            raise TypeError, "The class must define either the 'cppcode' " + \
                  "attribute or the 'eval' function"
        
        # If cppcode is provided send it to instant
        if 'cppcode' in dictionary:
            print "Calling instant to compile code and return class"
            cpp_BaseFunction = instant(dictionary['cppcode'])
            # Remove the cppcode attribute
            dictionary.pop('cppcode')
        # check type of eval
        elif not isinstance(dictionary['eval'],types.FunctionType):
            raise TypeError, "'eval' must be a 'function'"

        # Extend the bases with the cpp_BaseFunction and ffc_Function
        bases = (Function,cpp_BaseFunction,ffc_Function)

        # Defining the init function

        # If a user defined __init__ function is provided store it and call it
        # inside the created __init__ function.
        if '__init__' in dictionary:
            userdefined_init = dictionary['__init__']
            def __init__(self,V,*arg,**kwarg):
                print "Initializing %s"%name

                if not isinstance(V,FunctionSpace):
                    raise TypeError, "The first argument must be a 'FunctionSpace'"
                
                # Calling the user defined_init
                userdefined_init(self,V,*arg,**kwarg)
                
                self._V = V
                cpp_BaseFunction.__init__(self,V)
                ffc_Function.__init__(self,V.ffc_element())
                
        else:
            def __init__(self,V):
                print "Initializing %s"%name
                if not isinstance(V,FunctionSpace):
                    raise TypeError, "The argument must be a 'FunctionSpace'"
                self._V = V
                cpp_BaseFunction.__init__(self,V)
                ffc_Function.__init__(self,V.ffc_element())
        
        # Attach the init function to the dictionary
        dictionary['__init__'] = __init__

        # Check for the dim and rank attribute
        if not ('dim' in dictionary and 'rank' in dictionary):
            raise TypeError, "The class '%s' must implement both 'dim' and 'rank'"

        # Check the type of the dim attribute
        _dim = dictionary['dim']
        if not isinstance(_dim,(int,types.FunctionType)):
            raise "The attribute 'dim' must be a 'function' or an 'int'"
        
        if isinstance(_dim,int):
            # Check the passed integer
            if _dim not in [1,2,3]:
                raise ValueError, "'dim' must be 1, 2 or 3."
            
            # Define a dim function
            def dim(self):
                return _dim

            # Attach the dim function to dictionary
            dictionary['dim'] = dim

        # Check the type of the rank attribute
        _rank = dictionary['rank']
        if not isinstance(_rank,(int,types.FunctionType)):
            raise "The attribute 'rank' must be a 'function' or an 'int'"
        
        if isinstance(_rank,int):
            # Check the passed integer
            if _rank not in [0,1,2]:
                raise ValueError, "'rank' must be 1, 2 or 3."
            
            # Define a rank function
            def rank(self):
                return _rank
            
            # Attach the rank function to dictionary
            dictionary['rank'] = rank
        
        print "Bases on leaving:", bases
        print "Dictionary on leaving:",dictionary
        return type.__new__(cls, name, bases, dictionary)

class Function(object):
    __metaclass__ = FunctionCreator
    def __init__(self,V):
        if type(self) == Function:
            raise TypeError, "Can only instantiate subclasses of 'Function'"
        raise RuntimeError, "Do not initialize the 'Function' base class"
    
    def function_space(self):
        return self._V

print ""
if __name__ == '__main__':
    from dolfin import UnitSquare, FiniteElement
    # User example of how we could define user defined functions in python
    class MyCustomFunction(Function):
        def __init__(self,V,dummy):
            # Do not call Funcion.__init__(self,V) if you do a RunTime error is raised.
            print "Custom initalizing"
            self.dummy = dummy
        
        def eval(self,val,x):
            val[0] = 0.5
            val[1] = 0.5
    
        def dim(self):
            return 2
    
        def rank(self):
            return 1
    
    # Can also declare dim and rank as int attribute, see below
        
    print ""
    # Compiling a Function
    class MyCustomCompiledFunction(Function):
        cppcode = "Jada"
        dim  = 1
        rank = 0
    
    mesh = UnitSquare(10,10)
    element = FiniteElement('Lagrange','triangle',1)
    
    print ""
    V = FunctionSpace(mesh,element)
    
    print ""
    f0 = MyCustomFunction(V,"Dummy")
    
    print ""
    f1 = MyCustomCompiledFunction(V)
    
    try:
        f2 = Function(V)
    except Exception, e:
        print e