# ***************************************************************************
# *   Copyright (c) 2017 Johannes Hartung <j.hartung@gmx.net>               *
# *                                                                         *
# *   This program is free software; you can redistribute it and/or modify  *
# *   it under the terms of the GNU Lesser General Public License (LGPL)    *
# *   as published by the Free Software Foundation; either version 2 of     *
# *   the License, or (at your option) any later version.                   *
# *   for detail see the LICENCE text file.                                 *
# *                                                                         *
# *   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 Library General Public License for more details.                  *
# *                                                                         *
# *   You should have received a copy of the GNU Library General Public     *
# *   License along with this program; if not, write to the Free Software   *
# *   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  *
# *   USA                                                                   *
# *                                                                         *
# ***************************************************************************

__title__ = "FreeCAD FEM solver Fenics tools"
__author__ = "Johannes Hartung"
__url__ = "http://www.freecadweb.org"

## @package Fenics
#  \ingroup FEM

import FreeCAD

try:
    import fenics
except:
    FreeCAD.Console.PrintError("No Fenics modules found, please install them.")
else:
    import numpy as np

    class XDMFReader(object):
        """
        Reads XDMF file and provides unified interface for returning
        cell functions or facet functions.
        """
        def __init__(self, xdmffilename):
            """
            Sets filename and sets mesh instance to None.
            """
            self.xdmffilename = xdmffilename
            self.mesh = None

        def resetMesh(self):
            """
            Resets mesh instance to None.
            """
            self.mesh = None

        def readMesh(self):
            """
            If mesh instance is None, read mesh instance from file denoted
            by filename property.
            """
            # TODO: implement mesh read in for open file
            if self.mesh is None:
                xdmffile = fenics.XDMFFile(self.xdmffilename)
                self.mesh = fenics.Mesh()
                xdmffile.read(self.mesh)
                xdmffile.close()

        def readCellExpression(self, group_value_dict, value_type="scalar", overlap=lambda x: x[0], *args, **kwargs):
            """
            Reads cell expression and returns it.
            """
            value_type_dictionary = {
                "scalar": ScalarCellExpressionFromXDMF,
                "vector2d": Vector2DCellExpressionFromXDMF,
                "vector3d": Vector3DCellExpressionFromXDMF}

            self.readMesh()
            xdmffile = fenics.XDMFFile(self.xdmffilename)
            cf = value_type_dictionary[value_type.lower()](group_value_dict,
                                                           overlap=overlap,
                                                           *args, **kwargs)
            cf.init()
            for (key, value) in cf.group_value_dict.items():
                cf.markers[key] = fenics.MeshFunction("size_t", self.mesh, self.mesh.topology().dim())
                xdmffile.read(cf.markers[key], key)
                cf.dx[key] = fenics.Measure("dx", domain=self.mesh, subdomain_data=cf.markers[key])
            xdmffile.close()
            return cf

        def readFacetFunction(self, group_value_dict, *args, **kwargs):
            """
            Reads facet function and returns it.
            """
            self.readMesh()
            xdmffile = fenics.XDMFFile(self.xdmffilename)
            ff = FacetFunctionFromXDMF(group_value_dict, *args, **kwargs)
            ff.init()
            for (key, value) in ff.group_value_dict.items():
                ff.markers[key] = fenics.MeshFunction("size_t", self.mesh, self.mesh.topology().dim() - 1)
                xdmffile.read(ff.markers[key], key)
                ff.marked[key] = value.get("marked", 1)
                ff.ds[key] = fenics.Measure("ds", domain=self.mesh, subdomain_data=ff.markers[key])
                ff.bcs[key] = value
            xdmffile.close()
            return ff

    class CellExpressionFromXDMF(object):
        """
        Creates cell function expression from XDMF file.
        """
        def __init__(self, group_value_dict,
                     default=lambda x: 0.,
                     check_marked=(lambda x: x == 1), overlap=lambda x: x[0],
                     **kwargs):
            self.init()
            self.group_value_dict = group_value_dict
            self.check_marked = check_marked
            self.overlap = overlap
            self.default = default

        def init(self):
            self.markers = {}
            self.dx = {}

        def assign_values(self, values, to_assign):
            values[:] = to_assign

        def eval_cell_backend(self, values, x, cell):

            values_list = [func(x) for (key, func) in self.group_value_dict.items()
                           if self.check_marked(self.markers[key][cell.index])]
            return_value = self.overlap(values_list)

            if values_list:
                self.assign_values(values, return_value)
            else:
                self.assign_values(values, self.default(x))

            # TODO: python classes much slower than JIT compilation

    # ***********************************
    # * Sub classes due to value_shape method which is not of dynamical return type
    # * Also the assignment of values is to be done by reference. Therefore it has to be
    # * overloaded.
    # ***********************************

    class ScalarCellExpressionFromXDMF(fenics.Expression, CellExpressionFromXDMF):

        def __init__(self, group_value_dict,
                     default=lambda x: 0.,
                     check_marked=(lambda x: x == 1), overlap=lambda x: x[0], **kwargs):
            CellExpressionFromXDMF.__init__(self, group_value_dict,
                                            default=default,
                                            check_marked=check_marked,
                                            overlap=overlap)

        def eval_cell(self, values, x, cell):
            self.eval_cell_backend(values, x, cell)

        def value_shape(self):
            return ()

    class Vector3DCellExpressionFromXDMF(fenics.Expression, CellExpressionFromXDMF):

        def __init__(self, group_value_dict,
                     default=lambda x: np.zeros((3,)),
                     check_marked=(lambda x: x == 1), overlap=lambda x: x[0], **kwargs):
            CellExpressionFromXDMF.__init__(self, group_value_dict,
                                            default=default,
                                            check_marked=check_marked,
                                            overlap=overlap)

        def eval_cell(self, values, x, cell):
            self.eval_cell_backend(values, x, cell)

        def value_shape(self):
            return (3,)

    class Vector2DCellExpressionFromXDMF(fenics.Expression, CellExpressionFromXDMF):

        def __init__(self, group_value_dict,
                     default=lambda x: np.zeros((2,)),
                     check_marked=(lambda x: x == 1), overlap=lambda x: x[0], **kwargs):
            CellExpressionFromXDMF.__init__(self, group_value_dict,
                                            default=default,
                                            check_marked=check_marked,
                                            overlap=overlap)

        def eval_cell(self, values, x, cell):
            self.eval_cell_backend(values, x, cell)

        def value_shape(self):
            return (2,)

    class FacetFunctionFromXDMF(object):
        """
        Creates facet function from XDMF file.
        """
        def __init__(self, group_value_dict, *args, **kwargs):
            self.group_value_dict = group_value_dict
            self.init()

        def init(self):
            self.markers = {}
            self.marked = {}
            self.ds = {}
            self.bcs = {}

        def getDirichletBCs(self, vectorspace, *args, **kwargs):
            dbcs = []
            for (dict_key, dict_value) in self.bcs.items():
                if dict_value["type"] == "Dirichlet":
                    bc = fenics.DirichletBC(vectorspace, dict_value["value"], self.markers[dict_key], dict_value.get("marked", 1), *args, **kwargs)
                    dbcs.append(bc)
            return dbcs
        # TODO: write some functions to return integrals for Neumann and Robin
        # boundary conditions for the general case (i.e. vector, tensor)

##  @}
