#!/usr/bin/env python
#  -*- coding: utf-8 -*-
"""Functions to write mesh to different file formats.
Based on meshmagick <https://github.com/LHEEA/meshmagick> by François Rongère.
"""
# Copyright (C) 2017-2019 Matthieu Ancellin, based on the work of François Rongère
# See LICENSE file at <https://github.com/mancellin/capytaine>

import os
import time
import numpy as np

from capytaine.tools.optional_imports import import_optional_dependency


def write_mesh(filename, vertices, faces, file_format):
    """Driver function that writes every mesh file file_format known by meshmagick

    Parameters
    ----------
    filename: str
        name of the mesh file to be written on disk
    vertices: ndarray
        numpy array of the coordinates of the mesh's nodes
    faces: ndarray
        numpy array of the faces' nodes connectivities
    file_format: str
        file_format of the mesh defined in the extension_dict dictionary
    """

    if file_format not in extension_dict:
        raise IOError('Extension "%s" is not known' % file_format)

    writer = extension_dict[file_format]

    writer(filename, vertices, faces)


def write_DAT(filename, vertices, faces):
    """Writes .DAT file format for the DIODORE (PRINCIPA (c)) software.

    It also displays suggestions for inclusion into the .INP configuration
    file.

    Parameters
    ----------
    filename: str
        name of the mesh file to be written on disk
    vertices: ndarray
        numpy array of the coordinates of the mesh's nodes
    faces: ndarray
        numpy array of the faces' nodes connectivities
    """

    root_filename, ext = os.path.splitext(filename)
    filename = root_filename + ext.upper()
    ofile = open(filename, 'w')

    ofile.write('$\n$ Data for DIODORE input file : {0}\n'.format(root_filename.upper()))
    ofile.write('$ GENERATED BY MESHMAGICK ON {0}\n$\n'.format(time.strftime('%c')))

    ofile.write('$ NODE\n')
    vertex_block = \
        ''.join(
            (
                '\n'.join(
                    ''.join(
                        (
                            '{:8d}'.format(idx+1),
                            ''.join('{:13.5E}'.format(elt) for elt in node)
                        )
                    ) for (idx, node) in enumerate(vertices)
                ),

                '\n*RETURN\n'
            )
        )
    ofile.write(vertex_block)

    quad_block = '$\n$ ELEMENT,TYPE=Q4C000,ELSTRUCTURE={0}'.format(root_filename.upper())
    tri_block = '$\n$ ELEMENT,TYPE=T3C000,ELSTRUCTURE={0}'.format(root_filename.upper())
    nq = 0
    nt = 0
    for (idx, cell) in enumerate(faces+1):
        if cell[0] != cell[-1]:
            # quadrangle
            nq += 1
            quad_block = ''.join(
                (
                    quad_block,
                    '\n',
                    '{:8d}'.format(idx+1),
                    ''.join('{:8d}'.format(node_id) for node_id in cell)
                )
            )

        else:
            # Triangle
            nt += 1
            tri_block = ''.join(
                (
                    tri_block,
                    '\n',
                    '{:8d}'.format(idx+1),
                    ''.join('{:8d}'.format(node_id) for node_id in cell[:3])
                )
            )

    print('-------------------------------------------------')
    print('Suggestion for .inp DIODORE input file :')
    print('')
    print('*NODE,INPUT={0},FRAME=???'.format(root_filename))

    if nq > 0:
        quad_block = ''.join((quad_block, '\n*RETURN\n'))
        ofile.write(quad_block)
        print('*ELEMENT,TYPE=Q4C000,ELSTRUCTURE={0},INPUT={0}'.format(root_filename))
    if nt > 0:
        tri_block = ''.join((tri_block, '\n*RETURN\n'))
        ofile.write(tri_block)
        print('*ELEMENT,TYPE=T3C000,ELSTRUCTURE={0},INPUT={0}'.format(root_filename))

    print('')
    print('-------------------------------------------------')
    ofile.close()


def write_HST(filename, vertices, faces):
    """Writes .HST file format for the HYDROSTAR (Bureau Veritas (c)) software.

    Parameters
    ----------
    filename: str
        name of the mesh file to be written on disk
    vertices: ndarray
        numpy array of the coordinates of the mesh's nodes
    faces: ndarray
        numpy array of the faces' nodes connectivities
    """
    # TODO: allow many bodies

    ofile = open(filename, 'w')

    ofile.write(''.join((
        'PROJECT:\n',
        'USERS:   meshmagick\n\n'
        'NBODY   1\n'
        'RHO   1025.0\n'
        'GRAVITY   9.81\n\n'
    )))

    coordinates_block = ''.join((  # block
            'COORDINATES\n',
            '\n'.join(  # line
                ''.join(
                    (
                        '{:10d}'.format(idx+1),  # index
                        ''.join('{:16.6E}'.format(elt) for elt in node)  # node coordinates
                    )
                ) for (idx, node) in enumerate(vertices)
            ),
            '\nENDCOORDINATES\n\n'
    ))

    ofile.write(coordinates_block)

    cells_coordinates = ''.join((  # block
        'PANEL TYPE 0\n',
        '\n'.join(  # line
            ''.join(
                '{:10d}'.format(node_idx) for node_idx in cell
            ) for cell in faces + 1
        ),
        '\nENDPANEL\n\n'
    ))

    ofile.write(cells_coordinates)

    ofile.write('ENDFILE\n')

    ofile.close()


def write_TEC(filename, vertices, faces):
    """Writes .TEC file format for the TECPLOT (Tecplot (c)) visualisation software.

    Parameters
    ----------
    filename: str
        name of the mesh file to be written on disk
    vertices: ndarray
        numpy array of the coordinates of the mesh's nodes
    faces: ndarray
        numpy array of the faces' nodes connectivities
    """

    ofile = open(filename, 'w')

    nv = vertices.shape[0]
    nf = faces.shape[0]

    ofile.write('TITLE = \" THIS FILE WAS GENERATED BY MESHMAGICK - FICHIER : {} \" \n'.format(filename))

    ofile.write('VARIABLES = \"X\",\"Y\",\"Z\" \n')
    ofile.write('ZONE T=\"MESH\" \n')
    ofile.write('N={nv:10d} ,E={nf:10d} , F=FEPOINT, ET=QUADRILATERAL\n'.format(nv=nv, nf=nf))

    node_block = '\n'.join( # block
        ''.join(
            ''.join('{:16.6E}'.format(elt) for elt in node)
        ) for node in vertices
    ) + '\n'
    ofile.write(node_block)

    cells_block = '\n'.join(  # block
        ''.join(
            ''.join('{:10d}'.format(node_id) for node_id in cell)
        ) for cell in faces + 1
    ) + '\n'
    ofile.write(cells_block)

    ofile.close()

    return 1


def write_VTU(filename, vertices, faces):
    """Writes .vtu file format for the paraview (Kitware (c)) visualisation software.

    It relies on the VTK library for its writer. VTU files use the last XML file format of the VTK library.

    Parameters
    ----------
    filename: str
        name of the mesh file to be written on disk
    vertices: ndarray
        numpy array of the coordinates of the mesh's nodes
    faces: ndarray
        numpy array of the faces' nodes connectivities
    """
    vtk = import_optional_dependency("vtk")

    writer = vtk.vtkXMLUnstructuredGridWriter()
    writer.SetDataModeToAscii()
    writer.SetFileName(str(filename))

    unstructured_grid = _build_vtkUnstructuredGrid(vertices, faces)
    writer.SetInputData(unstructured_grid)
    writer.Write()


def write_VTP(filename, vertices, faces):
    """Writes .vtp file format for the Paraview (Kitware (c)) visualisation software.

    It relies on the VTK library for its writer. VTP files use the last XML file format of the VTK library and
    correspond to polydata.

    Parameters
    ----------
    filename: str
        name of the mesh file to be written on disk
    vertices: ndarray
        numpy array of the coordinates of the mesh's nodes
    faces: ndarray
        numpy array of the faces' nodes connectivities
    """

    vtk = import_optional_dependency("vtk")

    writer = vtk.vtkXMLPolyDataWriter()
    writer.SetDataModeToAscii()
    writer.SetFileName(str(filename))

    polydata = _build_vtkPolyData(vertices, faces)
    writer.SetInputData(polydata)
    writer.Write()


def write_VTK(filename, vertices, faces):
    """Writes .vtk file format for the Paraview (Kitware (c)) visualisation software.

    It relies on the VTK library for its writer. VTK files use the legagy ASCII file format of the VTK library.

    Parameters
    ----------
    filename: str
        name of the mesh file to be written on disk
    vertices: ndarray
        numpy array of the coordinates of the mesh's nodes
    faces: ndarray
        numpy array of the faces' nodes connectivities
    """

    nv = vertices.shape[0]
    nf = faces.shape[0]

    triangle_mask = (faces[:, 0] == faces[:, -1])
    quadrangles_mask = np.invert(triangle_mask)
    nb_triangles = len(np.where(triangle_mask)[0])
    nb_quandrangles = len(np.where(quadrangles_mask)[0])

    with open(filename, 'w') as f:

        f.write('# vtk DataFile Version 4.0\n')
        f.write('vtk file generated by meshmagick on %s\n' % time.strftime('%c'))
        f.write('ASCII\n')
        f.write('DATASET POLYDATA\n')
        f.write('POINTS %u float\n' % nv)

        for vertex in vertices:
            f.write('%f %f %f\n' % (vertex[0], vertex[1], vertex[2]))

        f.write('POLYGONS %u %u\n' % (nf, 4*nb_triangles+5*nb_quandrangles))

        for face in faces:
            if face[0] == face[-1]:  # Triangle
                f.write('3 %u %u %u\n' % (face[0], face[1], face[2]))
            else:  # Quadrangle
                f.write('4 %u %u %u %u\n' % (face[0], face[1], face[2], face[3]))


def _build_vtkUnstructuredGrid(vertices, faces):
    """Internal function that builds a VTK object for manipulation by the VTK library.

    Parameters
    ----------
    vertices: ndarray
        numpy array of the coordinates of the mesh's nodes
    faces: ndarray
        numpy array of the faces' nodes connectivities

    Returns
    -------
    vtkObject
    """

    vtk = import_optional_dependency("vtk")

    nv = max(np.shape(vertices))
    nf = max(np.shape(faces))

    vtk_mesh = vtk.vtkUnstructuredGrid()
    vtk_mesh.Allocate(nf, nf)

    # Building the vtkPoints data structure
    vtk_points = vtk.vtkPoints()
    vtk_points.SetNumberOfPoints(nv)
    for idx, vertex in enumerate(vertices):
        vtk_points.SetPoint(idx, vertex)

    vtk_mesh.SetPoints(vtk_points)  # Storing the points into vtk_mesh

    # Building the vtkCell data structure
    for cell in faces:
        if cell[-1] in cell[:-1]:
            vtk_cell = vtk.vtkTriangle()
            nc = 3
        else:
            # #print 'quadrangle'
            vtk_cell = vtk.vtkQuad()
            nc = 4

        for k in range(nc):
            vtk_cell.GetPointIds().SetId(k, cell[k])

        vtk_mesh.InsertNextCell(vtk_cell.GetCellType(), vtk_cell.GetPointIds())
    return vtk_mesh


def _build_vtkPolyData(vertices, faces):
    """Builds a vtkPolyData object from vertices and faces"""

    vtk = import_optional_dependency("vtk")

    # Create a vtkPoints object and store the points in it
    points = vtk.vtkPoints()
    for point in vertices:
        points.InsertNextPoint(point)

    # Create a vtkCellArray to store faces
    cell_array = vtk.vtkCellArray()
    for face_ids in faces:
        if face_ids[0] == face_ids[-1]:
            # Triangle
            curface = face_ids[:3]
            vtk_face = vtk.vtkTriangle()
        else:
            # Quadrangle
            curface = face_ids[:4]
            vtk_face = vtk.vtkQuad()

        for idx, id in enumerate(curface):
            vtk_face.GetPointIds().SetId(idx, id)

        cell_array.InsertNextCell(vtk_face)

    polydata_mesh = vtk.vtkPolyData()
    polydata_mesh.SetPoints(points)
    polydata_mesh.SetPolys(cell_array)

    return polydata_mesh


def write_NAT(filename, vertices, faces):
    """Writes .nat file format as defined into the load_NAT function.

    Parameters
    ----------
    filename: str
        name of the mesh file to be written on disk
    vertices: ndarray
        numpy array of the coordinates of the mesh's nodes
    faces: ndarray
        numpy array of the faces' nodes connectivities

    See Also
    --------
    load_NAT
    """

    ofile = open(filename, 'w')

    nv = max(np.shape(vertices))
    nf = max(np.shape(faces))

    ofile.write('%6u%6u\n' % (0, 0))  # lire les symmetries dans args...
    ofile.write('%6u%6u\n' % (nv, nf))
    for vertex in vertices:
        ofile.write('%15.6E%15.6E%15.6E\n' % (vertex[0], vertex[1], vertex[2]))
    for cell in faces+1:
        ofile.write('%10u%10u%10u%10u\n' % (cell[0], cell[1], cell[2], cell[3]))

    ofile.close()


def write_NEM(filename, vertices, faces):
    """Writes mesh files used by the Mesh tool included in Nemoh

    Parameters
    ----------
    filename : str
        name of the mesh file to be written on disk
    vertices: ndarray
        numpy array of the coordinates of the mesh's nodes
    faces: ndarray
        numpy array of the faces' nodes connectivities

    Note
    ----
    This file format is different from that used by Nemoh itself. It is only used by the Mesh tool.
    """
    ofile = open(filename, 'w')

    ofile.write('%u\n' % vertices.shape[0])
    ofile.write('%u\n' % faces.shape[0])

    for vertex in vertices:
        ofile.write('%15.6f\t%15.6f\t%15.6f\n' % (vertex[0], vertex[1], vertex[2]))

    for face in faces+1:
        ofile.write('%10u\t%10u\t%10u\t%10u\n' % (face[0], face[1], face[2], face[3]))

    ofile.close()


def write_GDF(filename, vertices, faces):
    """Writes .gdf file format for the WAMIT (Wamit INC. (c)) BEM software.

    Parameters
    ----------
    filename: str
        name of the mesh file to be written on disk
    vertices: ndarray
        numpy array of the coordinates of the mesh's nodes
    faces: ndarray
        numpy array of the faces' nodes connectivities
    """

    nf = max(np.shape(faces))

    ofile = open(filename, 'w')

    ofile.write('GDF file generated by meshmagick on %s\n' % time.strftime('%c'))

    ofile.write('%16.6f%16.6f\n' % (100.0, 9.81))
    ofile.write('%12u%12u\n' % (0, 1))  # TODO : mettre les symetries en argument
    ofile.write('%12u\n' % nf)

    for cell in faces:
        for k in range(4):
            cur_vertices = vertices[cell[k], :]
            ofile.write('%16.6E%16.6E%16.6E\n' % (cur_vertices[0], cur_vertices[1], cur_vertices[2]))

    ofile.close()


def write_MAR(filename, vertices, faces):
    """Writes mesh files to be used with Nemoh BEM software (Ecole Centrale de Nantes)

    Parameters
    ----------
    filename: str
        name of the mesh file to be written on disk
    vertices: ndarray
        numpy array of the coordinates of the mesh's nodes
    faces: ndarray
        numpy array of the faces' nodes connectivities
    """

    # TODO: detect symmetry in Oxz plane

    ofile = open(filename, 'w')

    ofile.write('{0:6d}{1:6d}\n'.format(2, 0))  # TODO : mettre les symetries en argument

    for (idx, vertex) in enumerate(vertices):
        ofile.write('{0:6d}{1:16.6f}{2:16.6f}{3:16.6f}\n'.format(idx+1, vertex[0], vertex[1], vertex[2]))

    ofile.write('{0:6d}{1:6d}{2:6d}{3:6d}{4:6d}\n'.format(0, 0, 0, 0, 0))

    cell_block = '\n'.join(
        ''.join('{0:10d}'.format(elt) for elt in cell)
        for cell in faces + 1
    ) + '\n'
    ofile.write(cell_block)
    ofile.write('%6u%6u%6u%6u\n' % (0, 0, 0, 0))

    ofile.close()

    print('WARNING: if you described only one part of the mesh using symmetry for Nemoh, you may manually modify the ' \
          'file header accordingly')


def write_RAD(filename, vertices, faces):
    raise NotImplementedError


def write_STL(filename, vertices, faces):
    """Writes .stl file format. It relies on the VTK library for its writer.

    Parameters
    ----------
    filename: str
        name of the mesh file to be written on disk
    vertices: ndarray
        numpy array of the coordinates of the mesh's nodes
    faces: ndarray
        numpy array of the faces' nodes connectivities
    """

    # TODO : replace this implementation by using the vtk functionalities

    # Triangulating quads
    t1 = (0, 1, 2)
    t2 = (0, 2, 3)

    quads_ids = np.where(faces[:, 0] != faces[:, -1])[0]

    new_faces = faces[quads_ids].copy()
    new_faces[:, :3] = new_faces[:, t1]
    new_faces[:, -1] = new_faces[:, 0]

    faces[quads_ids, :3] = faces[:, t2][quads_ids]
    faces[quads_ids, -1] = faces[quads_ids, 0]

    faces = np.concatenate((faces, new_faces))

    # Writing file
    ofile = open(filename, 'w')

    ofile.write('solid meshmagick\n')

    for face in faces:
        if face[0] != face[3]:
            raise RuntimeError("""Only full triangle meshes are accepted in STL files.
              Please consider using the --triangulate-quadrangles option (-tq) to
              perform a prior triangulation of the mesh""")

        # Computing normal
        v0 = vertices[face[0], :]
        v1 = vertices[face[1], :]
        v2 = vertices[face[2], :]

        n = np.cross(v1 - v0, v2 - v0)
        n /= np.linalg.norm(n)

        block_facet = ''.join(['  facet normal ', ''.join('%15.6e' % ni for ni in n) + '\n',
                               '    outer loop\n',
                               '      vertex', ''.join('%15.6e' % Vi for Vi in v0) + '\n',
                               '      vertex', ''.join('%15.6e' % Vi for Vi in v1) + '\n',
                               '      vertex', ''.join('%15.6e' % Vi for Vi in v2) + '\n',
                               '    endloop\n',
                               '  endfacet\n'])
        ofile.write(block_facet)
    ofile.write('endsolid meshmagick\n')
    ofile.close()


def write_INP(filename, vertices, faces):
    raise NotImplementedError('INP writer is not implementer yet')


def write_MSH(filename, vertices, faces):
    raise NotImplementedError('MSH writer is not implemented yet')


def write_MED(filename, vertices, faces):
    raise NotImplementedError('MED writer is not implemented yet')


def write_WRL(filename, vertices, faces):
    raise NotImplementedError('VRML writer is not implemented yet')


extension_dict = {  # keyword,  writer
    'mar': write_MAR,
    'nemoh': write_MAR,
    'wamit': write_GDF,
    'gdf': write_GDF,
    'diodore-inp': write_INP,
    'inp': write_INP,
    'diodore-dat': write_DAT,
    'hydrostar': write_HST,
    'hst': write_HST,
    'natural': write_NAT,
    'nat': write_NAT,
    'gmsh': write_MSH,
    'msh': write_MSH,
    'rad': write_RAD,
    'radioss': write_RAD,
    'stl': write_STL,
    'vtu': write_VTU,
    'vtp': write_VTP,
    'paraview-legacy': write_VTK,
    'vtk': write_VTK,
    'tecplot': write_TEC,
    'tec': write_TEC,
    'med': write_MED,
    'salome': write_MED,
    'vrml': write_WRL,
    'wrl': write_WRL,
    'nem': write_NEM,
    'nemoh_mesh': write_NEM
}
