"""Module containing useful plotting tools."""

from __future__ import annotations

from enum import Enum
import os
import platform
import subprocess
from subprocess import PIPE
from subprocess import Popen
from subprocess import TimeoutExpired
import sys

import numpy as np

import pyvista

from . import _vtk
from .colors import Color


class FONTS(Enum):
    """Font families available to PyVista."""

    arial = _vtk.VTK_ARIAL
    courier = _vtk.VTK_COURIER
    times = _vtk.VTK_TIMES


# Track render window support and plotting
SUPPORTS_OPENGL = None
SUPPORTS_PLOTTING = None


def supports_open_gl():
    """
    Return if the system supports OpenGL.

    This function checks if the system supports OpenGL by creating a VTK render
    window and querying its OpenGL support.

    Returns
    -------
    bool
        ``True`` if the system supports OpenGL, ``False`` otherwise.
    """
    global SUPPORTS_OPENGL
    if SUPPORTS_OPENGL is None:
        ren_win = _vtk.vtkRenderWindow()
        SUPPORTS_OPENGL = bool(ren_win.SupportsOpenGL())
    return SUPPORTS_OPENGL


def _system_supports_plotting():
    """Check if the environment supports plotting on Windows, Linux, or Mac OS.

    Returns
    -------
    system_supports_plotting : bool
        ``True`` when system supports plotting.

    """
    if os.environ.get('ALLOW_PLOTTING', '').lower() == 'true':
        return True

    # Windows case
    if os.name == 'nt':
        # actually have to check here.  Somewhat expensive.
        return supports_open_gl()

    # mac case
    if platform.system() == 'Darwin':
        # check if finder available
        proc = Popen(["pgrep", "-qx", "Finder"], stdout=PIPE, stderr=PIPE, encoding="utf8")
        try:
            proc.communicate(timeout=10)
        except TimeoutExpired:
            return False
        if proc.returncode == 0:
            return True

        # display variable set, likely available
        return 'DISPLAY' in os.environ

    # Linux case
    try:
        proc = Popen(["xset", "-q"], stdout=PIPE, stderr=PIPE, encoding="utf8")
        proc.communicate(timeout=10)
    except (OSError, TimeoutExpired):
        return False
    else:  # pragma: no cover
        return proc.returncode == 0


def system_supports_plotting():
    """Check if the environment supports plotting.

    Returns
    -------
    bool
        ``True`` when system supports plotting.

    """
    global SUPPORTS_PLOTTING
    if SUPPORTS_PLOTTING is None:
        SUPPORTS_PLOTTING = _system_supports_plotting()

    # always use the cached response
    return SUPPORTS_PLOTTING


def _update_axes_label_color(axes_actor, color=None):
    """Set the axes label color (internal helper)."""
    color = Color(color, default_color=pyvista.global_theme.font.color)
    if isinstance(axes_actor, _vtk.vtkAxesActor):
        prop_x = axes_actor.GetXAxisCaptionActor2D().GetCaptionTextProperty()
        prop_y = axes_actor.GetYAxisCaptionActor2D().GetCaptionTextProperty()
        prop_z = axes_actor.GetZAxisCaptionActor2D().GetCaptionTextProperty()
        for prop in [prop_x, prop_y, prop_z]:
            prop.SetColor(color.float_rgb)
            prop.SetShadow(False)
    elif isinstance(axes_actor, _vtk.vtkAnnotatedCubeActor):
        axes_actor.GetTextEdgesProperty().SetColor(color.float_rgb)


def create_axes_marker(
    label_color=None,
    x_color=None,
    y_color=None,
    z_color=None,
    xlabel='X',
    ylabel='Y',
    zlabel='Z',
    labels_off=False,
    line_width=2,
    cone_radius=0.4,
    shaft_length=0.8,
    tip_length=0.2,
    ambient=0.5,
    label_size=(0.25, 0.1),
):
    """Create an axis actor.

    Parameters
    ----------
    label_color : ColorLike, optional
        Color of the label text.

    x_color : ColorLike, optional
        Color of the x-axis text.

    y_color : ColorLike, optional
        Color of the y-axis text.

    z_color : ColorLike, optional
        Color of the z-axis text.

    xlabel : str, default: "X"
        Text used for the x-axis.

    ylabel : str, default: "Y"
        Text used for the y-axis.

    zlabel : str, default: "Z"
        Text used for the z-axis.

    labels_off : bool, default: False
        Enable or disable the text labels for the axes.

    line_width : float, default: 2
        The width of the marker lines.

    cone_radius : float, default: 0.4
        The radius of the axes arrow tips.

    shaft_length : float, default: 0.8
        The length of the axes arrow shafts.

    tip_length : float, default: 0.2
        Length of the tip.

    ambient : float, default: 0.5
        The ambient of the axes arrows. Value should be between 0 and 1.

    label_size : sequence[float], default: (0.25, 0.1)
        The width and height of the axes label actors. Values should be between
        0 and 1. For example ``(0.2, 0.1)``.

    Returns
    -------
    vtk.vtkAxesActor
        Axes actor.

    Examples
    --------
    Create the default axes marker.

    >>> import pyvista as pv
    >>> marker = pv.create_axes_marker()
    >>> pl = pv.Plotter()
    >>> _ = pl.add_actor(marker)
    >>> pl.show()

    Create an axes marker at the origin with custom colors and axis labels.

    >>> import pyvista as pv
    >>> marker = pv.create_axes_marker(
    ...     line_width=4,
    ...     ambient=0.0,
    ...     x_color="#378df0",
    ...     y_color="#ab2e5d",
    ...     z_color="#f7fb9a",
    ...     xlabel="X Axis",
    ...     ylabel="Y Axis",
    ...     zlabel="Z Axis",
    ...     label_size=(0.1, 0.1),
    ... )
    >>> pl = pv.Plotter()
    >>> _ = pl.add_actor(marker)
    >>> pl.show()

    """
    x_color = Color(x_color, default_color=pyvista.global_theme.axes.x_color)
    y_color = Color(y_color, default_color=pyvista.global_theme.axes.y_color)
    z_color = Color(z_color, default_color=pyvista.global_theme.axes.z_color)
    axes_actor = _vtk.vtkAxesActor()
    axes_actor.GetXAxisShaftProperty().SetColor(x_color.float_rgb)
    axes_actor.GetXAxisTipProperty().SetColor(x_color.float_rgb)
    axes_actor.GetYAxisShaftProperty().SetColor(y_color.float_rgb)
    axes_actor.GetYAxisTipProperty().SetColor(y_color.float_rgb)
    axes_actor.GetZAxisShaftProperty().SetColor(z_color.float_rgb)
    axes_actor.GetZAxisTipProperty().SetColor(z_color.float_rgb)
    # Set labels
    axes_actor.SetXAxisLabelText(xlabel)
    axes_actor.SetYAxisLabelText(ylabel)
    axes_actor.SetZAxisLabelText(zlabel)
    if labels_off:
        axes_actor.AxisLabelsOff()
    # Set Line width
    axes_actor.GetXAxisShaftProperty().SetLineWidth(line_width)
    axes_actor.GetYAxisShaftProperty().SetLineWidth(line_width)
    axes_actor.GetZAxisShaftProperty().SetLineWidth(line_width)

    axes_actor.SetConeRadius(cone_radius)
    axes_actor.SetNormalizedShaftLength([shaft_length] * 3)
    axes_actor.SetNormalizedTipLength([tip_length] * 3)
    axes_actor.GetXAxisShaftProperty().SetAmbient(ambient)
    axes_actor.GetYAxisShaftProperty().SetAmbient(ambient)
    axes_actor.GetZAxisShaftProperty().SetAmbient(ambient)
    axes_actor.GetXAxisTipProperty().SetAmbient(ambient)
    axes_actor.GetYAxisTipProperty().SetAmbient(ambient)
    axes_actor.GetZAxisTipProperty().SetAmbient(ambient)

    for label_actor in [
        axes_actor.GetXAxisCaptionActor2D(),
        axes_actor.GetYAxisCaptionActor2D(),
        axes_actor.GetZAxisCaptionActor2D(),
    ]:
        label_actor.SetWidth(label_size[0])
        label_actor.SetHeight(label_size[1])

    _update_axes_label_color(axes_actor, label_color)

    return axes_actor


def create_axes_orientation_box(
    line_width=1,
    text_scale=0.366667,
    edge_color='black',
    x_color=None,
    y_color=None,
    z_color=None,
    xlabel='X',
    ylabel='Y',
    zlabel='Z',
    x_face_color='red',
    y_face_color='green',
    z_face_color='blue',
    color_box=False,
    label_color=None,
    labels_off=False,
    opacity=0.5,
    show_text_edges=False,
):
    """Create a Box axes orientation widget with labels.

    Parameters
    ----------
    line_width : float, optional
        The width of the marker lines.

    text_scale : float, optional
        Size of the text relative to the faces.

    edge_color : ColorLike, optional
        Color of the edges.

    x_color : ColorLike, optional
        Color of the x-axis text.

    y_color : ColorLike, optional
        Color of the y-axis text.

    z_color : ColorLike, optional
        Color of the z-axis text.

    xlabel : str, optional
        Text used for the x-axis.

    ylabel : str, optional
        Text used for the y-axis.

    zlabel : str, optional
        Text used for the z-axis.

    x_face_color : ColorLike, optional
        Color used for the x-axis arrow.  Defaults to theme axes
        parameters.

    y_face_color : ColorLike, optional
        Color used for the y-axis arrow.  Defaults to theme axes
        parameters.

    z_face_color : ColorLike, optional
        Color used for the z-axis arrow.  Defaults to theme axes
        parameters.

    color_box : bool, optional
        Enable or disable the face colors.  Otherwise, box is white.

    label_color : ColorLike, optional
        Color of the labels.

    labels_off : bool, optional
        Enable or disable the text labels for the axes.

    opacity : float, optional
        Opacity in the range of ``[0, 1]`` of the orientation box.

    show_text_edges : bool, optional
        Enable or disable drawing the vector text edges.

    Returns
    -------
    vtk.vtkAnnotatedCubeActor
        Annotated cube actor.

    Examples
    --------
    Create and plot an orientation box

    >>> import pyvista as pv
    >>> actor = pv.create_axes_orientation_box(
    ...     line_width=1,
    ...     text_scale=0.53,
    ...     edge_color='black',
    ...     x_color='k',
    ...     y_color=None,
    ...     z_color=None,
    ...     xlabel='X',
    ...     ylabel='Y',
    ...     zlabel='Z',
    ...     color_box=False,
    ...     labels_off=False,
    ...     opacity=1.0,
    ... )
    >>> pl = pv.Plotter()
    >>> _ = pl.add_actor(actor)
    >>> pl.show()

    """
    x_color = Color(x_color, default_color=pyvista.global_theme.axes.x_color)
    y_color = Color(y_color, default_color=pyvista.global_theme.axes.y_color)
    z_color = Color(z_color, default_color=pyvista.global_theme.axes.z_color)
    edge_color = Color(edge_color, default_color=pyvista.global_theme.edge_color)
    x_face_color = Color(x_face_color)
    y_face_color = Color(y_face_color)
    z_face_color = Color(z_face_color)
    axes_actor = _vtk.vtkAnnotatedCubeActor()
    axes_actor.SetFaceTextScale(text_scale)
    if xlabel is not None:
        axes_actor.SetXPlusFaceText(f"+{xlabel}")
        axes_actor.SetXMinusFaceText(f"-{xlabel}")
    if ylabel is not None:
        axes_actor.SetYPlusFaceText(f"+{ylabel}")
        axes_actor.SetYMinusFaceText(f"-{ylabel}")
    if zlabel is not None:
        axes_actor.SetZPlusFaceText(f"+{zlabel}")
        axes_actor.SetZMinusFaceText(f"-{zlabel}")
    axes_actor.SetFaceTextVisibility(not labels_off)
    axes_actor.SetTextEdgesVisibility(show_text_edges)
    # https://github.com/pyvista/pyvista/pull/5382
    # axes_actor.GetTextEdgesProperty().SetColor(edge_color.float_rgb)
    axes_actor.GetTextEdgesProperty().SetLineWidth(line_width)
    axes_actor.GetXPlusFaceProperty().SetColor(x_color.float_rgb)
    axes_actor.GetXMinusFaceProperty().SetColor(x_color.float_rgb)
    axes_actor.GetYPlusFaceProperty().SetColor(y_color.float_rgb)
    axes_actor.GetYMinusFaceProperty().SetColor(y_color.float_rgb)
    axes_actor.GetZPlusFaceProperty().SetColor(z_color.float_rgb)
    axes_actor.GetZMinusFaceProperty().SetColor(z_color.float_rgb)

    axes_actor.GetCubeProperty().SetOpacity(opacity)
    axes_actor.GetCubeProperty().SetEdgeColor(edge_color.float_rgb)
    axes_actor.GetCubeProperty().SetEdgeVisibility(True)
    axes_actor.GetCubeProperty().BackfaceCullingOn()
    if opacity < 1.0:
        # Hide the text edges
        axes_actor.GetTextEdgesProperty().SetOpacity(0)

    if color_box:
        # Hide the cube so we can color each face
        axes_actor.GetCubeProperty().SetOpacity(0)
        axes_actor.GetCubeProperty().SetEdgeVisibility(False)

        cube = pyvista.Cube()
        cube.clear_data()  # remove normals
        face_colors = np.array(
            [
                x_face_color.int_rgb,
                x_face_color.int_rgb,
                y_face_color.int_rgb,
                y_face_color.int_rgb,
                z_face_color.int_rgb,
                z_face_color.int_rgb,
            ],
            np.uint8,
        )
        cube.cell_data['face_colors'] = face_colors

        cube_mapper = _vtk.vtkPolyDataMapper()
        cube_mapper.SetInputData(cube)
        cube_mapper.SetColorModeToDirectScalars()
        cube_mapper.Update()

        cube_actor = pyvista.Actor(mapper=cube_mapper)
        cube_actor.prop.culling = 'back'
        cube_actor.prop.opacity = opacity

        prop_assembly = _vtk.vtkPropAssembly()
        prop_assembly.AddPart(axes_actor)
        prop_assembly.AddPart(cube_actor)
        actor = prop_assembly
    else:
        actor = axes_actor

    _update_axes_label_color(actor, label_color)

    return actor


def create_north_arrow():
    """Create a north arrow mesh.

    .. versionadded:: 0.44.0

    Returns
    -------
    pyvista.PolyData
        North arrow mesh.

    """
    points = np.array(
        [
            [0.0, 5.0, 0.0],
            [-2.0, 0.0, 0.0],
            [0.0, 1.5, 0.0],
            [2.0, 0.0, 0.0],
            [0.0, 5.0, 1.0],
            [-2.0, 0.0, 1.0],
            [0.0, 1.5, 1.0],
            [2.0, 0.0, 1.0],
        ],
    )
    faces = np.array(
        [
            4,
            3,
            7,
            4,
            0,
            4,
            2,
            6,
            7,
            3,
            4,
            1,
            5,
            6,
            2,
            4,
            0,
            4,
            5,
            1,
            4,
            0,
            1,
            2,
            3,
            4,
            4,
            7,
            6,
            5,
        ],
    )
    return pyvista.PolyData(points, faces)


def normalize(x, minimum=None, maximum=None):
    """
    Normalize the given value between [minimum, maximum].

    Parameters
    ----------
    x : numpy.ndarray
        The array of values to normalize.
    minimum : float, optional
        The minimum value to which ``x`` should be normalized. If not specified,
        the minimum value in ``x`` will be used.
    maximum : float, optional
        The maximum value to which ``x`` should be normalized. If not specified,
        the maximum value in ``x`` will be used.

    Returns
    -------
    numpy.ndarray
        The normalized array of values, where the values are scaled to the
        range ``[minimum, maximum]``.
    """
    if minimum is None:
        minimum = np.nanmin(x)
    if maximum is None:
        maximum = np.nanmax(x)
    return (x - minimum) / (maximum - minimum)


def opacity_transfer_function(mapping, n_colors, interpolate=True, kind='quadratic'):
    """Get the opacity transfer function for a mapping.

    These values will map on to a scalar bar range and thus the number of
    colors (``n_colors``) must correspond to the number of colors in the color
    mapping that these opacities are associated to.

    If interpolating, ``scipy.interpolate.interp1d`` is used if available,
    otherwise ``np.interp`` is used. The ``kind`` argument controls the kind of
    interpolation for ``interp1d``.

    This returns the opacity range from 0 to 255, where 0 is totally
    transparent and 255 is totally opaque.

    The equation to create the sigmoid mapping is: ``1 / (1 + exp(-x))`` where
    ``x`` is the range from ``-a`` to ``+a`` and ``a`` is the value given in
    the ``mapping`` string. Default is ``a=10`` for 'sigmoid' mapping.

    Parameters
    ----------
    mapping : list(float) or str
        The opacity mapping to use. Can be a ``str`` name of a predefined
        mapping including ``'linear'``, ``'geom'``, ``'sigmoid'``,
        ``'sigmoid_1-10,15,20'``, and ``foreground``. Append an ``'_r'`` to any
        of those names (except ``foreground``) to reverse that mapping.
        The mapping can also be a custom user-defined array/list of values
        that will be interpolated across the ``n_color`` range.

    n_colors : int
        The number of colors that the opacities must be mapped to.

    interpolate : bool
        Flag on whether or not to interpolate the opacity mapping for all
        colors.

    kind : str
        The interpolation kind if ``interpolate`` is ``True`` and ``scipy``
        is available. If ``scipy`` is not available, linear interpolation
        is always used. Options are:

        - ``'linear'``
        - ``'nearest'``
        - ``'zero'``
        - ``'slinear'``
        - ``'quadratic'``
        - ``'cubic'``
        - ``'previous'``
        - ``'next'``

    Returns
    -------
    numpy.ndarray
        Array of ``numpy.uint8`` values ``n_colors`` long containing the
        [0-255] opacity mapping values.

    Examples
    --------
    >>> import pyvista as pv
    >>> # Fetch the `sigmoid` mapping between 0 and 255
    >>> tf = pv.opacity_transfer_function("sigmoid", 256)
    >>> # Fetch the `geom_r` mapping between 0 and 1
    >>> tf = (
    ...     pv.opacity_transfer_function("geom_r", 256).astype(float)
    ...     / 255.0
    ... )
    >>> # Interpolate a user defined opacity mapping
    >>> opacity = [0, 0.2, 0.9, 0.6, 0.3]
    >>> tf = pv.opacity_transfer_function(opacity, 256)

    """
    sigmoid = lambda x: np.array(1 / (1 + np.exp(-x)) * 255, dtype=np.uint8)
    transfer_func = {
        'linear': np.linspace(0, 255, n_colors, dtype=np.uint8),
        'geom': np.geomspace(1e-6, 255, n_colors, dtype=np.uint8),
        'geom_r': np.geomspace(255, 1e-6, n_colors, dtype=np.uint8),
        'sigmoid': sigmoid(np.linspace(-10.0, 10.0, n_colors)),
        'sigmoid_1': sigmoid(np.linspace(-1.0, 1.0, n_colors)),
        'sigmoid_2': sigmoid(np.linspace(-2.0, 2.0, n_colors)),
        'sigmoid_3': sigmoid(np.linspace(-3.0, 3.0, n_colors)),
        'sigmoid_4': sigmoid(np.linspace(-4.0, 4.0, n_colors)),
        'sigmoid_5': sigmoid(np.linspace(-5.0, 5.0, n_colors)),
        'sigmoid_6': sigmoid(np.linspace(-6.0, 6.0, n_colors)),
        'sigmoid_7': sigmoid(np.linspace(-7.0, 7.0, n_colors)),
        'sigmoid_8': sigmoid(np.linspace(-8.0, 8.0, n_colors)),
        'sigmoid_9': sigmoid(np.linspace(-9.0, 9.0, n_colors)),
        'sigmoid_10': sigmoid(np.linspace(-10.0, 10.0, n_colors)),
        'sigmoid_15': sigmoid(np.linspace(-15.0, 15.0, n_colors)),
        'sigmoid_20': sigmoid(np.linspace(-20.0, 20.0, n_colors)),
        'foreground': np.hstack((0, [255] * (n_colors - 1))).astype(np.uint8),
    }
    transfer_func['linear_r'] = transfer_func['linear'][::-1]
    transfer_func['sigmoid_r'] = transfer_func['sigmoid'][::-1]
    for i in range(3, 11):
        k = f'sigmoid_{i}'
        rk = f'{k}_r'
        transfer_func[rk] = transfer_func[k][::-1]
    if isinstance(mapping, str):
        try:
            return transfer_func[mapping]
        except KeyError:
            raise ValueError(
                f'Opacity transfer function ({mapping}) unknown. '
                f'Valid options: {list(transfer_func.keys())}',
            ) from None
    elif isinstance(mapping, (np.ndarray, list, tuple)):
        mapping = np.array(mapping)
        if mapping.size == n_colors:
            # User could pass transfer function ready for lookup table
            pass
        elif mapping.size < n_colors:
            # User pass custom transfer function to be linearly interpolated
            if np.max(mapping) > 1.0 or np.min(mapping) < 0.0:
                mapping = normalize(mapping)
            # Interpolate transfer function to match lookup table
            xo = np.linspace(0, n_colors, len(mapping), dtype=np.int_)
            xx = np.linspace(0, n_colors, n_colors, dtype=np.int_)
            try:
                if not interpolate:
                    raise ValueError('No interpolation.')
                # Use a quadratic interp if scipy is available
                from scipy.interpolate import interp1d

                # quadratic has best/smoothest results
                f = interp1d(xo, mapping, kind=kind)
                vals = f(xx)
                vals[vals < 0] = 0.0
                vals[vals > 1.0] = 1.0
                mapping = (vals * 255.0).astype(np.uint8)

            except (ImportError, ValueError):
                # Otherwise use simple linear interp
                mapping = (np.interp(xx, xo, mapping) * 255).astype(np.uint8)
        else:
            raise RuntimeError(
                f'Transfer function cannot have more values than `n_colors`. This has {mapping.size} elements',
            )
        return mapping
    raise TypeError(f'Transfer function type ({type(mapping)}) not understood')


def parse_font_family(font_family: str) -> int:
    """
    Check and validate the given font family name.

    Parameters
    ----------
    font_family : str
        Font family name to validate. Must be one of the font names defined in
        the ``FONTS`` enum class.

    Returns
    -------
    int
        Corresponding integer value of the valid font family name in the
        ``FONTS`` enum class.

    Raises
    ------
    ValueError
        If the font_family is not one of the defined font names in the ``FONTS``
        enum class.
    """
    font_family = font_family.lower()
    fonts = [font.name for font in FONTS]
    if font_family not in fonts:
        raise ValueError(f'Font must one of the following:\n{", ".join(fonts)}')
    return FONTS[font_family].value


def check_matplotlib_vtk_compatibility():
    """
    Check if VTK and Matplotlib versions are compatible for MathText rendering.

    This function is primarily geared towards checking if MathText rendering is
    supported with the given versions of VTK and Matplotlib. It follows the
    version constraints:

    * VTK <= 9.2.2 requires Matplotlib < 3.6
    * VTK > 9.2.2 requires Matplotlib >= 3.6

    Other version combinations of VTK and Matplotlib will work without
    errors, but some features (like MathText/LaTeX rendering) may
    silently fail.

    Returns
    -------
    bool
        True if the versions of VTK and Matplotlib are compatible for MathText
        rendering, False otherwise.

    Raises
    ------
    RuntimeError
        If the versions of VTK and Matplotlib cannot be checked.

    """
    import matplotlib as mpl

    mpl_vers = tuple(map(int, mpl.__version__.split('.')[:2]))
    if pyvista.vtk_version_info <= (9, 2, 2):
        return not mpl_vers >= (3, 6)
    elif pyvista.vtk_version_info > (9, 2, 2):
        return mpl_vers >= (3, 6)
    raise RuntimeError('Uncheckable versions.')  # pragma: no cover


def check_math_text_support():
    """Check if MathText and LaTeX symbols are supported.

    Returns
    -------
    bool
        ``True`` if both MathText and LaTeX symbols are supported, ``False``
        otherwise.
    """
    # Something seriously sketchy is happening with this VTK code
    # It seems to hijack stdout and stderr?
    # See https://github.com/pyvista/pyvista/issues/4732
    # This is a hack to get around that by executing the code in a subprocess
    # and capturing the output:
    # _vtk.vtkMathTextFreeTypeTextRenderer().MathTextIsSupported()
    _cmd = "import vtk;print(vtk.vtkMathTextFreeTypeTextRenderer().MathTextIsSupported());"
    proc = subprocess.run([sys.executable, '-c', _cmd], check=False, capture_output=True)
    math_text_support = False if proc.returncode else proc.stdout.decode().strip() == 'True'
    return math_text_support and check_matplotlib_vtk_compatibility()