# -*- coding: utf-8 -*-
"""This module defines a Form class that wraps FFC forms and UFC forms
into a cpp.Form (dolfin::Form)."""

# Copyright (C) 2008-2015 Johan Hake
#
# This file is part of DOLFIN.
#
# DOLFIN is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# DOLFIN 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 Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with DOLFIN. If not, see <http://www.gnu.org/licenses/>.
#
# Modified by Anders Logg, 2008-2011
# Modified by Marie E. Rognes, 2011
# Modified by Martin Sandve Alnæs, 2014-2015

# Import SWIG-generated extension module (DOLFIN C++)
import dolfin.cpp as cpp

import ufl
from dolfin.compilemodules.jit import jit

from dolfin.cpp import deprecation
from dolfin.cpp import MeshFunctionSizet

__all__ = ["Form"]

# Note that we need to store _compiled_form and _compiled_coefficients
# to prevent Python from garbage-collecting these while still in use.
# FIXME: Figure out how to solve this with shared_ptr

function_space_error = "Error while extracting test and/or trial spaces. "


class Form(cpp.Form):
    def __init__(self, form,
                 form_compiler_parameters=None,
                 function_spaces=None):
        "Create JIT-compiled form from any given form (compiled or not)."

        # Developer note: The function_spaces argument is used to
        # handle forms defined on multimesh function spaces. These are
        # really defined over a standard function space (the first of
        # the spaces) then later recreated on each function space part
        # and added together to form a multimesh form. This happens in
        # the function assemble_multimesh.

        # Check form argument
        if not isinstance(form, ufl.Form):
            cpp.dolfin_error("form.py",
                             "creating dolfin.Form",
                             "Expected a ufl.Form.")

        # Cutoff for better error message than ffc would give
        if not form.empty() and not form.ufl_domains():
            cpp.dolfin_error("form.py",
                             "creating dolfin.Form",
                             "The ufl.Form does not include any reference to a mesh.")

        # Extract subdomain data.
        # Until the assembler and ufc has multimesh support,
        # we impose the limitation of a single domain at this point.
        sd = form.subdomain_data()
        if len(sd) != 1:
            cpp.dolfin_error("form.py",
                             "creating dolfin.Form",
                             "Expecting subdomain data for exactly one domain"
                             "(got {}).".format(len(sd)))
        self.subdomains, = list(sd.values())  # Assuming single domain
        domain, = list(sd.keys())  # Assuming single domain
        mesh = domain.ufl_cargo()
        if isinstance(mesh, cpp.MultiMesh):
            mesh = mesh.part(0)

        # Having a mesh in the form is now a strict requirement
        if mesh is None:
            cpp.dolfin_error("form.py",
                             "creating dolfin.Form",
                             "Expecting to find a Mesh in the form.")

        # Jit the module, this will take some time...
        jit_result = jit(form, form_compiler_parameters,
                         mpi_comm=mesh.mpi_comm())
        if jit_result is None:
            cpp.dolfin_error("form.py",
                             "creating dolfin.Form",
                             "jit failure.")
        self._compiled_form, module, prefix = jit_result

        # Extract function spaces of form arguments
        if function_spaces is None:
            self.function_spaces = [func.function_space() for func
                                    in form.arguments()]
        else:
            self.function_spaces = function_spaces

        # Type checking argument function_spaces
        if not all(isinstance(fs, cpp.FunctionSpace)
                   for fs in self.function_spaces):
            raise ValueError(function_space_error +
                             " Invalid type of test spaces.")

        # Initialize base class
        cpp.Form.__init__(self, self._compiled_form,
                          self.function_spaces)

        # Type checking argument function_spaces
        #r = self._compiled_form.rank()
        #if len(self.function_spaces) != r:
        #    raise ValueError(function_space_error +
        #                     " Wrong number of test spaces (should be %d)." % r)

        # Extract coefficients from form (pass only the ones that
        # the compiled form actually uses to the assembler)
        original_coefficients = form.coefficients()
        self.coefficients = []
        for i in range(self.num_coefficients()):
            j = self.original_coefficient_position(i)
            self.coefficients.append(original_coefficients[j])

        # Type checking coefficients
        if not all( isinstance(c, (cpp.GenericFunction, cpp.MultiMeshFunction))
                   for c in self.coefficients):
            # Developer note:
            # The form accepts a MultiMeshFunction but does not set the
            # correct coefficients. This is done in assemble_multimesh
            # at the moment
            coefficient_error = "Error while extracting coefficients. "
            raise TypeError(coefficient_error +
                            "Either provide a dict of cpp.GenericFunctions, " +
                            "or use Function to define your form.")

        for i in range(self.num_coefficients()):
            if isinstance(self.coefficients[i], cpp.GenericFunction):
                self.set_coefficient(i, self.coefficients[i])

        # Attach mesh (because function spaces and coefficients may be
        # empty lists)
        if not function_spaces:
            self.set_mesh(mesh)
        # Type checking subdomain data
        # Delete None entries
        for k in list(self.subdomains.keys()):
            if self.subdomains[k] is None:
                del self.subdomains[k]
        # Check that we have no additional subdomain data (user
        # misspelling)
        # TODO: Get from ufl list?
        integral_types = ("cell", "exterior_facet", "interior_facet", "vertex")
        additional_keys = set(self.subdomains.keys()) - set(integral_types)
        if additional_keys:
            cpp.warning("Invalid keys in subdomains: %s" % additional_keys)
        # Check that we have only MeshFunctions
        for data in list(self.subdomains.values()):
            if not (data is None or isinstance(data, MeshFunctionSizet)):
                cpp.warning("Invalid subdomain data type %s, expecting a MeshFunction" % type(data))

        # Attach subdomains to C++ Form if we have them
        subdomains = self.subdomains.get("cell")
        if subdomains is not None:
            self.set_cell_domains(subdomains)
        subdomains = self.subdomains.get("exterior_facet")
        if subdomains is not None:
            self.set_exterior_facet_domains(subdomains)
        subdomains = self.subdomains.get("interior_facet")
        if subdomains is not None:
            self.set_interior_facet_domains(subdomains)
        subdomains = self.subdomains.get("vertex")
        if subdomains is not None:
            self.set_vertex_domains(subdomains)
