# ***************************************************************************
# *                                                                         *
# *   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                                                                   *
# *                                                                         *
# ***************************************************************************
from __future__ import print_function

from .importToolsFem import \
    get_FemMeshObjectDimension,\
    get_FemMeshObjectElementTypes,\
    get_MaxDimElementFromList,\
    get_FemMeshObjectOrder,\
    get_FemMeshObjectMeshGroups
from xml.etree import ElementTree as ET  # parsing xml files and exporting
import numpy as np


__title__ = "FreeCAD Fenics XDMF mesh writer"
__author__ = "Johannes Hartung"
__url__ = "http://www.freecadweb.org"

## @package exportFenicsXDMF
#  \ingroup FEM
#  \brief FreeCAD Fenics Mesh XDMF writer for FEM workbench

ENCODING_ASCII = 'ASCII'
ENCODING_HDF5 = 'HDF5'

FreeCAD_Group_Dimensions = {
    "Vertex": 0,
    "Edge": 1,
    "Face": 2,
    "Volume": 3
}

FreeCAD_to_Fenics_XDMF_dict = {
    ("Node", 1): ("polyvertex", 1),
    ("Edge", 1): ("polyline", 2),
    ("Edge", 2): ("edge_3", 3),
    ("Triangle", 1): ("triangle", 3),
    ("Triangle", 2): ("tri_6", 6),
    ("Tetra", 1): ("tetrahedron", 4),
    ("Tetra", 2): ("tet_10", 10)
}

# we need numpy functions to later access and process large data sets in a fast manner
# also the hd5 support works better together with numpy


def numpy_array_to_str(npa):
    res = ""
    dt = str(npa.dtype)
    if 'int' in dt:
        res = "\n".join([" ".join([("%d" % s) for s in a]) for a in npa.tolist()])
    elif 'float' in dt:
        res = "\n".join([" ".join([("%3.6f" % s) for s in a]) for a in npa.tolist()])
    return res


def points_to_numpy(pts, dim=3):
    return np.array([[p.x, p.y, p.z] for p in pts])[:, :dim]


def tuples_to_numpy(tpls, numbers_per_line):
    return np.array([list(t) for t in tpls])[:, :numbers_per_line]


def write_fenics_mesh_points_xdmf(fem_mesh_obj, geometrynode, encoding=ENCODING_ASCII):
    """
        Writes either into hdf5 file or into open mesh file
    """

    numnodes = fem_mesh_obj.FemMesh.NodeCount

    dim = get_MaxDimElementFromList(get_FemMeshObjectElementTypes(fem_mesh_obj))[2]
    effective_dim = dim
    if dim <= 2:
        effective_dim = 2  # effective dim is 2 for dim==1
        geometrynode.set("GeometryType", "XY")
    elif dim == 3:
        geometrynode.set("GeometryType", "XYZ")

    recalc_nodes_ind_dict = {}

    if encoding == ENCODING_ASCII:
        dataitem = ET.SubElement(geometrynode, "DataItem", Dimensions="%d %d" % (numnodes, effective_dim), Format="XML")
        nodes = []
        for (ind, (key, node)) in enumerate(list(fem_mesh_obj.FemMesh.Nodes.items())):
            nodes.append(node)
            recalc_nodes_ind_dict[key] = ind

        dataitem.text = numpy_array_to_str(points_to_numpy(nodes, dim=effective_dim))
    elif encoding == ENCODING_HDF5:
        pass

    return recalc_nodes_ind_dict


def write_fenics_mesh_codim_xdmf(fem_mesh_obj,
                                 topologynode,
                                 nodes_dict,
                                 codim=0,
                                 encoding=ENCODING_ASCII):
    mesh_dimension = get_FemMeshObjectDimension(fem_mesh_obj)

    element_types = get_FemMeshObjectElementTypes(fem_mesh_obj, remove_zero_element_entries=True)
    element_order = get_FemMeshObjectOrder(fem_mesh_obj)
    # we get all elements from mesh to decide which one to write by selection of codim
    # nodeindices = [(nodes_dict[ind] for ind in fem_mesh_obj.FemMesh.getElementNodes(fc_volume_ind)) for (fen_ind, fc_volume_ind) in enumerate(fc_cells)]
    writeout_element_dimension = mesh_dimension - codim

    (num_topo, name_topo, dim_topo) = (0, "", 0)
    for (num, name, dim) in element_types:
        if writeout_element_dimension == dim:
            (num_topo, name_topo, dim_topo) = (num, name, dim)

    (topology_type, nodes_per_element) = FreeCAD_to_Fenics_XDMF_dict[(name_topo, element_order)]

    topologynode.set("TopologyType", topology_type)
    topologynode.set("NumberOfElements", str(num_topo))
    topologynode.set("NodesPerElement", str(nodes_per_element))

    if dim_topo == 3:
        fc_topo = fem_mesh_obj.FemMesh.Volumes
    elif dim_topo == 2:
        fc_topo = fem_mesh_obj.FemMesh.Faces
    elif dim_topo == 1:
        fc_topo = fem_mesh_obj.FemMesh.Edges
    elif dim_topo == 0:
        fc_topo = fem_mesh_obj.FemMesh.Nodes
    else:
        fc_topo = []
        print("Dimension of mesh incompatible with export XDMF function: %d" % (dim_topo,))

    nodeindices = [(nodes_dict[ind] for ind in fem_mesh_obj.FemMesh.getElementNodes(fc_topo_ind)) for (fen_ind, fc_topo_ind) in enumerate(fc_topo)]

    if encoding == ENCODING_ASCII:
        dataitem = ET.SubElement(topologynode, "DataItem", NumberType="UInt", Dimensions="%d %d" % (num_topo, nodes_per_element), Format="XML")
        dataitem.text = numpy_array_to_str(tuples_to_numpy(nodeindices, nodes_per_element))
    elif encoding == ENCODING_HDF5:
        pass

    return fc_topo


def write_fenics_mesh_scalar_cellfunctions(name, cell_array, attributenode, encoding=ENCODING_ASCII):
    attributenode.set("AttributeType", "Scalar")
    attributenode.set("Center", "Cell")
    attributenode.set("Name", name)

    (num_cells, num_dims) = np.shape(cell_array)

    if encoding == ENCODING_ASCII:
        dataitem = ET.SubElement(attributenode, "DataItem", Dimensions="%d %d" % (num_cells, num_dims), Format="XML")
        dataitem.text = numpy_array_to_str(cell_array)
    elif encoding == ENCODING_HDF5:
        pass


"""
Example: mesh with two topologies and one mesh function for the facet one

<?xml version="1.0"?>
<!DOCTYPE Xdmf SYSTEM "Xdmf.dtd" []>
<Xdmf Version="3.0" xmlns:xi="http://www.w3.org/2001/XInclude">
  <Domain>
    <Grid Name="mesh" GridType="Uniform">
      <Topology NumberOfElements="162" TopologyType="Tetrahedron" NodesPerElement="4">
        <DataItem Dimensions="162 4" NumberType="UInt" Format="XML">0 1 5 21
...
        </DataItem>
      </Topology>
      <Geometry GeometryType="XYZ">
        <DataItem Dimensions="64 3" Format="XML">0 0 0
...
        </DataItem>
      </Geometry>
    </Grid>
    <Grid Name="mesh" GridType="Uniform">
      <Topology NumberOfElements="378" TopologyType="Triangle" NodesPerElement="3">
        <DataItem Dimensions="378 3" NumberType="UInt" Format="XML">0 1 5
...
        </DataItem>
      </Topology>
      <Geometry Reference="XML">/Xdmf/Domain/Grid/Geometry</Geometry>
      <Attribute Name="f" AttributeType="Scalar" Center="Cell">
        <DataItem Dimensions="378 1" Format="XML">3
...
        </DataItem>
      </Attribute>
    </Grid>
  </Domain>
</Xdmf>


"""


def write_fenics_mesh_xdmf(fem_mesh_obj, outputfile, group_values_dict={}, encoding=ENCODING_ASCII):
    """
        For the export of xdmf.
    """

    FreeCAD_to_Fenics_dict = {
        "Triangle": "triangle",
        "Tetra": "tetrahedron",
        "Hexa": "hexahedron",
        "Edge": "interval",
        "Node": "point",
        "Quadrangle": "quadrilateral",

        "Polygon": "unknown", "Polyhedron": "unknown",
        "Prism": "unknown", "Pyramid": "unknown",
    }

    print("Converting " + fem_mesh_obj.Label + " to fenics XDMF File")
    print("Dimension of mesh: %d" % (get_FemMeshObjectDimension(fem_mesh_obj),))

    elements_in_mesh = get_FemMeshObjectElementTypes(fem_mesh_obj)
    print("Elements appearing in mesh: %s" % (str(elements_in_mesh),))
    celltype_in_mesh = get_MaxDimElementFromList(elements_in_mesh)
    (num_cells, cellname_fc, dim_cell) = celltype_in_mesh
    cellname_fenics = FreeCAD_to_Fenics_dict[cellname_fc]
    print("Celltype in mesh -> %s and its Fenics dolfin name: %s" % (str(celltype_in_mesh), cellname_fenics))

    root = ET.Element("Xdmf", version="3.0")
    domain = ET.SubElement(root, "Domain")
    base_grid = ET.SubElement(domain, "Grid", Name="base_mesh", GridType="Uniform")
    base_topology = ET.SubElement(base_grid, "Topology")
    base_geometry = ET.SubElement(base_grid, "Geometry")

    # TODO: for the general mesh: write out topology and geometry in grid node
    # TODO: for every marked group write own grid node with topology (ref if cells)
    #       geometry ref, attribute

    #####################################
    # write base topo and geometry
    nodes_dict = write_fenics_mesh_points_xdmf(fem_mesh_obj, base_geometry, encoding=encoding)
    write_fenics_mesh_codim_xdmf(fem_mesh_obj, base_topology, nodes_dict, codim=0, encoding=encoding)
    #####################################

    fem_mesh = fem_mesh_obj.FemMesh
    gmshgroups = get_FemMeshObjectMeshGroups(fem_mesh_obj)

    if gmshgroups is not ():
        print('found mesh groups')

    for g in gmshgroups:
        mesh_function_type = fem_mesh.getGroupElementType(g)
        mesh_function_codim = dim_cell - FreeCAD_Group_Dimensions[mesh_function_type]
        mesh_function_name = fem_mesh.getGroupName(g)

        print('group id: %d (label: %s) with element type %s and codim %d'
              % (g, mesh_function_name, mesh_function_type, mesh_function_codim))

        mesh_function_grid = ET.SubElement(domain, "Grid", Name=mesh_function_name + "_mesh", GridType="Uniform")
        mesh_function_topology = ET.SubElement(mesh_function_grid, "Topology")

        mesh_function_topology_description = write_fenics_mesh_codim_xdmf(fem_mesh_obj,
                                                                          mesh_function_topology,
                                                                          nodes_dict,
                                                                          codim=mesh_function_codim, encoding=encoding)

        mesh_function_geometry = ET.SubElement(mesh_function_grid, "Geometry", Reference="XML")
        mesh_function_geometry.text = "/Xdmf/Domain/Grid/Geometry"
        mesh_function_attribute = ET.SubElement(mesh_function_grid, "Attribute")

        elem_dict = {}
        (elem_mark_group, elem_mark_default) = group_values_dict.get(g, (1, 0))

        # TODO: is it better to save all groups each at once or collect all codim equal
        # groups to put them into one function?
        # TODO: nevertheless there has to be a dialog which fixes the default value and the mark value

        for e in fem_mesh.getGroupElements(g):
            elem_dict[e] = elem_mark_group

        val_array = np.array([elem_dict.get(e, elem_mark_default) for e in mesh_function_topology_description])
        topo_array = np.vstack((val_array,)).T
        write_fenics_mesh_scalar_cellfunctions(mesh_function_name,
                                               topo_array,
                                               mesh_function_attribute, encoding=ENCODING_ASCII)

    # TODO: improve cell functions support

    fp = open(outputfile, "wb")
    fp.write(b'''<?xml version="1.0"?>\n<!DOCTYPE Xdmf SYSTEM "Xdmf.dtd" []>\n''')
    fp.write(ET.tostring(root))
    # xml core functionality does not support pretty printing
    # so the output file looks quite ugly
    fp.close()
