# Copyright (C) 2011-2015 Claas Abert
#
# This file is part of magnum.fe. 
#
# magnum.fe 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.
#
# magnum.fe 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 magnum.fe. If not, see <http://www.gnu.org/licenses/>.
# 
# Last modified by Claas Abert, 2015-06-29

from __future__ import absolute_import 

from dolfin import Function, MeshFunction, Mesh, SubMesh, CellFunction, FunctionSpaceBase, interpolate
import numpy as np

__all__ = ["WrappedMesh"]

class WrappedMesh(Mesh):

  def __init__(self, *args, **kwargs):
    """
    This class represents a mesh with a certain submesh and defines
    methods for the fast interpolation between the two meshes.
    """
    super(WrappedMesh, self).__init__(*args, **kwargs)

  def _init(self, mesh_with_shell):
    """
    Initialize the mesh instance.

    *Arguments*
      mesh_with_shell (:class:`dolfin.Mesh`)
        The super mesh
    """
    self.with_shell = mesh_with_shell

    # Cache DOF mappings
    # The signature of the element is used as key
    self._mappings = {}

  def cut(self, f, **kwargs):
    """
    Takes a function defined on the super mesh and returns a truncated
    function defined on the sub mesh.

    *Arguments*
      f (:class:`dolfin.Function`)
        The function on the super mesh.

    *Returns*
      :class:`dolfin.Function`
        The function on the sub mesh.
    """

    mapping = self._get_mapping(f.function_space())
    result = Function(mapping['Vsub'])
    result.vector()[:] = f.vector().array()[mapping['map']]
    result.rename(f.name(), f.label())
    return result

  def expand(self, f, target = None):
    """
    Takes a function defined on the sub mesh and returns a function
    defined on the super mesh with unknown values set to zero.

    *Arguments*
      f (:class:`dolfin.Function`)
        The function on the sub mesh.

    *Returns*
      The function on the super mesh.
    """

    mapping = self._get_mapping(f.function_space())
    if target is None:
      target = Function(mapping['Vsuper'])
      target.rename(f.name(), f.label())
    target.vector()[mapping['map']] = f.vector().array()
    return target

  def _get_mapping(self, V):
    element = V.ufl_element()
    key = V.element().signature()

    if not self._mappings.has_key(key):
      Vsub   = FunctionSpaceBase(self, element)
      Vsuper = FunctionSpaceBase(self.with_shell, element)

      fsuper = Function(Vsuper)
      fsuper.vector()[:] = np.arange(fsuper.vector().size(), dtype=float)
      fsub = interpolate(fsuper, Vsub)

      self._mappings[key] = {
        'map':    np.round(fsub.vector().array()).astype(np.uint64),
        'Vsub':   Vsub,
        'Vsuper': Vsuper
      }

    return self._mappings[key]

  @staticmethod
  def create(mesh, domain_ids, invert=False):
    """
    Creates a wrapped mesh from a super mesh for a given collection
    of domain IDs.

    *Arguments*
      mesh (:class:`dolfin.Mesh`)
        The mesh.
      domain_ids (:class:`[int]`)
        List of domain IDs
      invert (:class:`bool`)
        Invert list of domain IDs

    *Returns*
      :class:`WrappedMesh`
        The wrapped mesh
    """
    if invert or isinstance(domain_ids, list) or isinstance(domain_ids, tuple):
      if isinstance(domain_ids, int): domain_ids = (domain_ids,)

      subdomains = MeshFunction('size_t', mesh, 3, mesh.domains())
      combined_subdomains = CellFunction("size_t", mesh, 0)
      for domain_id in domain_ids:
        combined_subdomains.array()[subdomains.array() == domain_id] = 1

      submesh = SubMesh(mesh, combined_subdomains, 0 if invert else 1)
    else:
      submesh = SubMesh(mesh, domain_ids)

    submesh.__class__  = WrappedMesh
    submesh._init(mesh)

    return submesh