# -*- coding: utf-8 -*-
"""
This module contains functionality for function spaces in particular
discrete function spaces defined over meshes in terms of finite
elements.
"""

# 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/>.

__all__ = ["FunctionSpace", "VectorFunctionSpace", "TensorFunctionSpace",
           "MultiMeshFunctionSpace"]

from six import string_types
import numpy

# Import UFL and SWIG-generated extension module (DOLFIN C++)
import ufl
import dolfin.cpp as cpp
from dolfin.compilemodules.jit import jit


def _get_common_mesh_and_constrained_domain(spaces):
    mesh = spaces[0].mesh()
    constrained_domain = spaces[0].dofmap().constrained_domain

    # Check that subspaces actually share the mesh
    if not all(V.mesh().id() == mesh.id() for V in spaces):
        cpp.dolfin_error("functionspace.py",
                         "create function space",
                         "Nonmatching meshes for function space: " \
                         + str([V.mesh() for V in spaces]))

    # Also check that all spaces share same constrained_domain map we
    # pass on
    ncd = sum(0 if V.dofmap().constrained_domain is None else 1 for V in spaces)
    if ncd == 0:
        # No constrained domains, nothing to check
        pass
    elif ncd == len(spaces):
        pass
        # Error checking here is not possible at the moment because
        # SubDomain is not a Variable and thus has no .id(), and the
        # swig layer gives us a different python object for the same
        # constrained_domain
        # All constrained domains should be the same
        #for V in spaces:
        #    cd2 = V.dofmap().constrained_domain
        #    if not cd2 is constrained_domain or cd2.id() == constrained_domain.id():
        #        cpp.dolfin_error("functionspace.py",
        #                         "create function space",
        #                         "Nonmatching constrained_domain for function space: " \
        #                         + str([V.dofmap().constrained_domain for V in spaces]))
    else:
        # Not allowing some to have constrained domain and some not to.
        cpp.dolfin_error("functionspace.py",
                         "create function space",
                         "Nonmatching constrained_domain for function space: " \
                         + str([V.dofmap().constrained_domain for V in spaces]))

    return mesh, constrained_domain


def _compile_dolfin_element(element, mesh, constrained_domain=None):
    "Instantiate DOLFIN FiniteElement and DofMap from ufl element."

    # JIT-compile element to get ufc_element and ufc_dofmap
    ufc_element, ufc_dofmap = jit(element, mpi_comm=mesh.mpi_comm())

    dolfin_element = cpp.FiniteElement(ufc_element)

    if constrained_domain is None:
        dolfin_dofmap  = cpp.DofMap(ufc_dofmap, mesh)
    else:
        dolfin_dofmap  = cpp.DofMap(ufc_dofmap, mesh, constrained_domain)

    return dolfin_element, dolfin_dofmap


class FunctionSpace(ufl.FunctionSpace, cpp.FunctionSpace):
    "Base class for all function spaces."

    def __init__(self, *args, **kwargs):
        """Create finite element function space.

        **Overloaded versions**

        * FunctionSpace(mesh, element, constrained_domain=None)

          Create function space on given mesh for given finite element

        * FunctionSpace(mesh, family, degree, form_degree=None,
                        constrained_domain=None, restriction=None)

          Convenience function for specifying element by string, degree, etc.

        * FunctionSpace(cppV)

          Internal library function for wrapping :py:class`cpp.FunctionSpace
          <dolfin.cpp.FunctionSpace>` into :py:class:`FunctionSpace
          <dolfin.functions.functionspace.FunctionSpace>`. No copying is done.

        *Arguments*
            mesh (:py:class:`Mesh <dolfin.cpp.Mesh>`)
                the mesh
            element (:py:class:`ufl.FiniteElement <ufl.FiniteElementBase>`)
                the element
            family (string)
                specification of the element family, see below for
                alternatives.
            degree (int)
                the degree of the element.
            form_degree (int)
                the form degree (FEEC notation, used when field is
                viewed as k-form)
            constrained_domain
                constrained subdomain with map function.
            restriction
                restriction of the element (e.g. to cell facets).

        Which families and degrees that are supported is determined by the
        form compiler used to generate the element, but typical families
        include

        =================================  =========
        Name                               Usage
        =================================  =========
        Argyris*                           "ARG"
        Arnold-Winther*                    "AW"
        Brezzi-Douglas-Fortin-Marini*      "BDFM"
        Brezzi-Douglas-Marini              "BDM"
        Bubble                             "B"
        Crouzeix-Raviart                   "CR"
        Discontinuous Lagrange             "DG"
        Discontinuous Raviart-Thomas       "DRT"
        Hermite*                           "HER"
        Lagrange                           "CG"
        Mardal-Tai-Winther*                "MTW"
        Morley*                            "MOR"
        Nedelec 1st kind H(curl)           "N1curl"
        Nedelec 2nd kind H(curl)           "N2curl"
        Quadrature                         "Quadrature"
        Raviart-Thomas                     "RT"
        Real                               "R"
        =================================  =========

        *only partly supported.

        *Examples of usage*
            To define a discrete function space over e.g. the unit square:

            .. code-block:: python

                mesh = UnitSquare(32,32)
                V = FunctionSpace(mesh, "CG", 1)

            Here, ``"CG"`` stands for Continuous Galerkin, implying the
            standard Lagrange family of elements. Instead of ``"CG"``, we
            could have written ``"Lagrange"``. With degree 1, we get the
            linear Lagrange element. Other examples include:

            .. code-block:: python

                # Discontinuous element of degree 0
                V = FunctionSpace(mesh, "DG", 0)

                # Brezzi-Douglas-Marini element of degree 2
                W = FunctionSpace(mesh, "BDM", 2)

                # Real element with one global degree of freedom
                R = FunctionSpace(mesh, "R", 0)

        """
        if len(args) == 1:
            self._init_from_cpp(*args, **kwargs)
        else:
            if len(args) == 0 or not isinstance(args[0], cpp.Mesh):
                cpp.dolfin_error("functionspace.py",
                                 "create function space",
                                 "Illegal argument, not a mesh: "
                                 + str(args[0]))
            elif len(args) == 2:
                self._init_from_ufl(*args, **kwargs)
            else:
                self._init_convenience(*args, **kwargs)

    def _init_from_ufl(self, mesh, element, constrained_domain=None):
        if not isinstance(element, ufl.FiniteElementBase):
            cpp.dolfin_error("functionspace.py",
                             "create function space",
                             "Illegal argument, not a finite element: "
                             + str(element))
        if constrained_domain is not None:
            if not isinstance(constrained_domain, cpp.SubDomain):
                cpp.dolfin_error("functionspace.py",
                                 "create function space",
                                 "Illegal argument, not a subdomain: "
                                 + str(constrained_domain))

        # Initialize the ufl.FunctionSpace first to check for good meaning
        ufl.FunctionSpace.__init__(self, mesh.ufl_domain(), element)

        # Compile element and dofmap and construct corresponding dolfin objects
        dolfin_element, dolfin_dofmap = _compile_dolfin_element(element, mesh, constrained_domain=constrained_domain)

        # Initialize the cpp.FunctionSpace
        cpp.FunctionSpace.__init__(self, mesh, dolfin_element, dolfin_dofmap)

        # Store reference to constrained domain to avoid possible SWIG
        # director memory error (see
        # https://bitbucket.org/fenics-project/dolfin/issue/71)
        # TODO: Make constrained_domain private with getter? Attribute access breaks interface consistency.
        self.constrained_domain = constrained_domain

    def _init_convenience(self, mesh, family, degree, form_degree=None,
                          constrained_domain=None, restriction=None):
        # Check arguments
        if not isinstance(family, string_types):
            cpp.dolfin_error("functionspace.py",
                             "create function space",
                             "Illegal argument for finite element family, not a string: " + str(family))
        if not isinstance(degree, int):
            cpp.dolfin_error("functionspace.py",
                             "create function space",
                             "Illegal argument for degree, not an integer: "
                             + str(degree))

        # Create UFL element
        element = ufl.FiniteElement(family, mesh.ufl_cell(), degree,
                                    form_degree=form_degree)
        if restriction is not None:
            element = element[restriction]

        self._init_from_ufl(mesh, element,
                            constrained_domain=constrained_domain)

    def _init_from_cpp(self, cppV, **kwargs):
        if not isinstance(cppV, cpp.FunctionSpace):
            cpp.dolfin_error("functionspace.py",
                             "create function space",
                             "Illegal argument for C++ function space, "
                             "not a cpp.FunctionSpace: " + str(cppV))
        # We don't want to support copy construction. This would indicate
        # internal defficiency in the library
        if isinstance(cppV, FunctionSpace):
            cpp.dolfin_error("functionspace.py",
                             "create function space",
                             "Illegal argument for C++ function space, "
                             "should not be functions.functionspace.FunctionSpace: " + str(cppV))
        if len(kwargs) > 0:
            cpp.dolfin_error("functionspace.py",
                             "create function space",
                             "Illegal arguments, did not expect C++ "
                             "function space and **kwargs: " + str(kwargs))

        # Assign all the members (including 'this' pointer to SWIG wraper)
        # NOTE: This in fact performs assignment of C++ context
        self.__dict__ = cppV.__dict__

        # Reconstruct UFL element from signature
        ufl_element = eval(self.element().signature(), ufl.__dict__)

        # Get mesh
        ufl_domain = self.mesh().ufl_domain()

        # Initialize the ufl.FunctionSpace (not calling cpp.Function.__init__)
        ufl.FunctionSpace.__init__(self, ufl_domain, ufl_element)

    def ufl_function_space(self):
        return self

    def ufl_cell(self):
        "Return the UFL cell."
        return self.ufl_domain().ufl_cell()

    def dolfin_element(self):
        "Return the DOLFIN element."
        return self.element()

    def __str__(self):
        "Pretty-print."
        return "<Function space of dimension %d (%s)>" % \
               (self.dofmap().global_dimension(), str(self.ufl_element()))

    def __eq__(self, other):
        "Comparison for equality."
        return ufl.FunctionSpace.__eq__(self, other) \
           and cpp.FunctionSpace.__eq__(self, other)

    def __ne__(self, other):
        "Comparison for inequality."
        return ufl.FunctionSpace.__ne__(self, other) \
            or cpp.FunctionSpace.__ne__(self, other)

    def num_sub_spaces(self):
        "Return the number of sub spaces"
        return self.dolfin_element().num_sub_elements()

    def sub(self, i):
        "Return the i-th sub space"
        # FIXME: Should we have a more extensive check other than
        # whats includedin the cpp code?
        if not isinstance(i, int):
            raise TypeError("expected an int for 'i'")
        if self.num_sub_spaces() == 1:
            raise ValueError("no SubSpaces to extract")
        if i >= self.num_sub_spaces():
            raise ValueError("Can only extract SubSpaces with i = 0 ... %d" % \
                  (self.num_sub_spaces() - 1))
        assert hasattr(self.ufl_element(), "sub_elements")

        # Extend with the python layer
        return FunctionSpace(cpp.FunctionSpace.sub(self, i))

    def extract_sub_space(self, component):
        """
        Extract subspace for component

        *Arguments*
            component (numpy.array(uint))
               The component.

        *Returns*
            _FunctionSpace_
                The subspace.
        """
        # Transform the argument to a NumPy array
        if not hasattr(component, "__len__"):
            cpp.dolfin_error("functionspace.py",
                             "extracting sub spaces",
                             "Expected a component which is iterable")
        component = numpy.asarray(component, dtype=numpy.uintp)

        # Get the cpp version of the FunctionSpace
        cpp_space = cpp.FunctionSpace.extract_sub_space(self, component)

        # Extend with the python layer
        return FunctionSpace(cpp_space)

    def split(self):
        """Split a mixed functionspace into its sub spaces"""
        return [self.sub(i) for i in range(self.num_sub_spaces())]

    def collapse(self, collapsed_dofs=False):
        """
        Collapse a subspace and return a new function space and a map
        from new to old dofs

        *Arguments*
            collapsed_dofs (bool)
                Return the map from new to old dofs

       *Returns*
           _FunctionSpace_
                The new function space.
           dict
                The map from new to old dofs (optional)
        """
        # Get the cpp version of the FunctionSpace
        cpp_space, dofs = cpp.FunctionSpace.collapse(self)

        # Extend with the python layer
        V = FunctionSpace(cpp_space)

        if collapsed_dofs:
            return V, dofs
        else:
            return V


def VectorFunctionSpace(mesh, family, degree, dim=None, form_degree=None,
                        constrained_domain=None, restriction=None):
    """Create vector-valued finite element function space.

    Use VectorFunctionSpace if the unknown is a vector field,
    instead of a :py:class:`FunctionSpace
    <dolfin.functions.functionspace.FunctionSpace>` object for
    scalar fields.

    *Arguments*
        mesh (:py:class:`Mesh <dolfin.cpp.Mesh>`)
            the mesh
        family (string)
            a string specifying the element family, see
            :py:class:`FunctionSpace
            <dolfin.functions.functionspace.FunctionSpace>` for
            alternatives.
        degree (int)
            the (polynomial) degree of the element.
        dim (int)
            an optional argument specifying the number of components.
        form_degree (int)
            an optional argument specifying the degree of the
            k-form (used for FEEC notation)

    If the dim argument is not provided, the dimension will be
    deduced from the dimension of the mesh.

    *Example of usage*

        .. code-block:: python

            V = VectorFunctionSpace(mesh, "CG", 1)

    """
    # Create element
    element = ufl.VectorElement(family, mesh.ufl_cell(), degree, dim=dim,
                                form_degree=form_degree)
    if restriction is not None:
        element = element[restriction]

    return FunctionSpace(mesh, element,
                         constrained_domain=constrained_domain)


def TensorFunctionSpace(mesh, family, degree, shape=None, symmetry=None,
                        constrained_domain=None, restriction=None):
    """Create tensor-valued finite element function space.

    *Arguments*
        mesh
            a :py:class:`Mesh <dolfin.cpp.Mesh>`.
        family
            a string specifying the element family,
            see :py:class:`FunctionSpace
            <dolfin.functions.functionspace.FunctionSpace>`
            for alternatives.
        degree
            the degree of the element.
        shape
            an optional argument specifying the shape of the tensor.
        symmetry
            optional argument specifying whether the tensor is symmetric.

    If the shape argument is not provided, the dimension will be
    deduced from the dimension of the mesh.

    *Example of usage*

        .. code-block:: python

            V = TensorFunctionSpace(mesh, "CG", 1)

    """
    # Create element
    element = ufl.TensorElement(family, mesh.ufl_cell(), degree, shape,
                                symmetry)
    if restriction is not None:
        element = element[restriction]

    return FunctionSpace(mesh, element,
                         constrained_domain=constrained_domain)


def MultiMeshFunctionSpace(multimesh, family, degree=None):
    """Create multimesh finite element function space.

    *Arguments*
        multimesh
            a :py:class:`MultiMesh <dolfin.cpp.MultiMesh>`.
        family
            a string specifying the element family,
            see :py:class:`FunctionSpace
            <dolfin.functions.functionspace.FunctionSpace>`
            for alternatives.

            This argument may also be a `FiniteElement`, in
            which case the `degree` argument should not be
            specified.
        degree
            the degree of the element.

    *Example of usage*

        .. code-block:: python

            V = MultiMeshFunctionSpace(mesh, "CG", 1)

            element = FiniteElement("Lagrange", triangle, 1)
            V = MultiMeshFunctionSpace(mesh, element)
    """

    # Check arguments
    if not isinstance(multimesh, cpp.MultiMesh):
        cpp.dolfin_error("functionspace.py",
                         "create multimesh function space",
                         "Illegal argument, not a multimesh: " + str(multimesh))

    # Create element if not supplied
    if isinstance(family, ufl.FiniteElementBase):
        element = family
    else:
        mesh = multimesh.part(0)
        element = ufl.FiniteElement(family, mesh.ufl_cell(), degree)

    # Create and add individual function spaces
    V = cpp.MultiMeshFunctionSpace(multimesh)
    V_parts = []
    for part in range(multimesh.num_parts()):
        V_part = FunctionSpace(multimesh.part(part), element)
        V_parts.append(V_part)
        V.add(V_part)

    # Build multimesh function space
    V.build()

    # Store full function spaces
    V._parts = V_parts

    return V
