"""Module for converting :class:`pymech.core.HexaData` objects to vtk"""
import os
from itertools import product
from pathlib import Path
from warnings import warn

import numpy as np

from .log import logger

logger.warning(
    "The module pymech.vtksuite is experimental in nature and "
    "may have some rough edges. The functions can also change in the future."
)

try:
    from tvtk.api import tvtk, write_data
except ImportError:
    logger.warning("To use VTK functions,\n    pip install mayavi")


__all__ = ("exa2vtk", "writevtk")


# ==============================================================================
def exa2vtk(field, downsample=False):
    warn(
        (
            "Function pymech.vtksuite.exa2vtk is now pymech.vtksuite.hexa2vtk. "
            "This module is kept for backwards "
            "compatibility and would disappear in version 2.0.0"
        ),
        DeprecationWarning,
    )
    return hexa2vtk(field, downsample)


def hexa2vtk(field, downsample=False):
    """A function for converting :class:`pymech.core.HexaData` to `Traited VTK`_ dataset. The
    returned dataset can be manipulated with libraries which accept a VTK
    object, for example Mayavi_.

    .. _Traited VTK: https://docs.enthought.com/mayavi/tvtk/README.html

    .. todo::

        Try https://github.com/pyvista/pyvista-xarray

    Example
    -------
    This also requires you to have a GUI toolkit installed: either PyQt4,
    PySide, PySide2, PyQt5 or wxPython.

    .. code-block:: python

       import pymech as pm
       from pymech.vtksuite import exa2vtk
       from mayavi import mlab

       field = pm.readnek("tests/nek/channel3D_0.f00001")
       dataset = exa2vtk(field)
       mlab.pipeline.add_dataset(dataset)

    Instead of MayaVi_ you could use also use something high-level like PyVista_
    to wrap the underlying VTK object and later visualize them.

    .. code-block:: python

        import pyvista as pv
        dataset = pv.wrap(dataset._vtk_obj)
        dataset.plot()

    .. _MayaVi: https://docs.enthought.com/mayavi/mayavi/mlab.html
    .. _PyVista: https://docs.pyvista.org/getting-started/index.html

    Parameters
    ----------
    field : :class:`pymech.core.HexaData`
            a dataset in nekdata format
    downsample : bool
            flag T/F

    Returns
    -------
    dataset : tvtk.tvtk_classes.unstructured_grid.UnstructuredGrid
            a VTK dataset
    """
    #
    if downsample:
        ixs = field.lr1[0] - 1
        iys = field.lr1[1] - 1
        izs = max(field.lr1[2] - 1, 1)
    else:
        ixs = 1
        iys = 1
        izs = 1
    #
    iix = range(0, field.lr1[0], ixs)
    nix = len(iix)
    iiy = range(0, field.lr1[1], iys)
    niy = len(iiy)
    iiz = range(0, field.lr1[2], izs)
    niz = len(iiz)
    #
    nppel = nix * niy * niz
    nepel = (nix - 1) * (niy - 1) * max((niz - 1), 1)
    nel = field.nel * nepel
    #
    if field.ndim == 3:
        nvert = 8
        cellType = tvtk.Hexahedron().cell_type
    else:
        nvert = 4
        cellType = tvtk.Quad().cell_type
    #
    ct = np.array(nel * [cellType])
    of = np.arange(0, nvert * nel, nvert)

    ce = np.zeros(nel * (nvert + 1))
    ce[np.arange(nel) * (nvert + 1)] = nvert

    if field.var[0] != 0:
        r = np.zeros((nvert * nel, 3))
    if field.var[1] != 0:
        v = np.zeros((nvert * nel, 3))
    if field.var[2] == 1:
        p = np.zeros(nvert * nel)
    if field.var[3] == 1:
        T = np.zeros(nvert * nel)
    if field.var[4] != 0:
        S = np.zeros((nvert * nel, field.var[4]))
    #
    ice = -(nvert + 1)

    for iel in range(field.nel):
        for (iz, ez), (iy, ey), (ix, ex) in product(
            enumerate(iiz), enumerate(iiy), enumerate(iix)
        ):
            iarray = iel * nppel + ix + iy * nix + iz * (nix * niy)

            # Downsample copy into a column vector
            if field.var[0] == 3:
                r[iarray, :] = field.elem[iel].pos[:, ez, ey, ex]
            if field.var[1] == 3:
                v[iarray, :] = field.elem[iel].vel[:, ez, ey, ex]
            if field.var[2] == 1:
                p[iarray] = field.elem[iel].pres[:, ez, ey, ex]
            if field.var[3] == 1:
                T[iarray] = field.elem[iel].temp[:, ez, ey, ex]
            if field.var[4] != 0:
                S[iarray, :] = field.elem[iel].scal[:, ez, ey, ex]
        if field.var[0] == 3:
            for iz, iy, ix in product(
                range(max(niz - 1, 1)), range(niy - 1), range(nix - 1)
            ):
                ice = ice + nvert + 1
                for face in range(field.ndim - 1):
                    cell_id = iel * nppel + ix + iy * nix + (iz + face) * nix * niy

                    ce[ice + face * 4 + 1] = cell_id
                    ce[ice + face * 4 + 2] = cell_id + 1
                    ce[ice + face * 4 + 3] = cell_id + nix + 1
                    ce[ice + face * 4 + 4] = cell_id + nix

    # create the array of cells
    ca = tvtk.CellArray()
    ca.set_cells(nel, ce)
    # create the unstructured dataset
    dataset = tvtk.UnstructuredGrid(points=r)
    # set the cell types
    dataset.set_cells(ct, of, ca)
    # set the data
    idata = 0
    if field.var[1] != 0:
        dataset.point_data.vectors = v
        dataset.point_data.vectors.name = "vel"
        idata += 1
    if field.var[2] == 1:
        dataset.point_data.scalars = p
        dataset.point_data.scalars.name = "pres"
        idata += 1
    if field.var[3] == 1:
        dataset.point_data.add_array(T)
        dataset.point_data.get_array(idata).name = "temp"
        idata += 1
    if field.var[4] != 0:
        for ii in range(field.var[4]):
            dataset.point_data.add_array(S[:, ii])
            dataset.point_data.get_array(ii + idata).name = "scal_%d" % (ii + 1)
    #
    dataset.point_data.update()
    #
    return dataset


def writevtk(fname, data):
    """A function for writing binary data in the XML VTK format

    Parameters
    ----------
    fname : str
            file name
    data : :class:`pymech.core.HexaData`
            data organised after reading a file

    """
    ext = ".vtp"
    fname = Path(fname)
    if fname.suffix != ext:
        logger.info(f"Renaming {fname} with extension .vtp")
        fname = fname.with_suffix(ext)

    vtk_dataset = hexa2vtk(data)
    write_data(vtk_dataset, os.fspath(fname))
