#
# Copyright 2018-2025 the orix developers
#
# This file is part of orix.
#
# orix is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# orix 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with orix. If not, see <http://www.gnu.org/licenses/>.
#

"""Private tools for coloring crystal directions.

These functions are adapted from MTEX.
"""

import matplotlib.colors as mcolors
import numpy as np

from orix.quaternion.rotation import Rotation
from orix.vector.fundamental_sector import FundamentalSector
from orix.vector.vector3d import Vector3d


def polar_coordinates_in_sector(
    sector: FundamentalSector, v: Vector3d
) -> tuple[np.ndarray, np.ndarray]:
    r"""Calculate the polar coordinates of crystal direction(s)
    relative to the (bary)center of a fundamental sector of a Laue
    group.

    Parameters
    ----------
    sector
        A fundamental sector of a Laue group, described by a set of
        normal vectors, a center, and vertices.
    v
        Crystal direction(s) to get polar coordinates for.

    Returns
    -------
    azimuth
        Azimuthal polar coordinate(s).
    polar
        Polar polar coordinate(s).

    Notes
    -----
    This procedure is adopted from MTEX' :code:`polarCoordinates`
    function, which implements the coloring described in section 2.4 in
    :cite:`nolze2016orientation` (see Fig. 4 in that reference).

    The azimuthal coordinate is the angle to the barycenter relative to
    some fixed vertex of the sector, here chosen as the north pole
    [001]. The polar coordinate is the distance to the barycenter. The
    barycenter is the center of the fundamental sector.
    """
    center = sector.center.unit
    n_vertices = sector.vertices.size
    v = v.unit

    # Azimuthal coordinate
    if n_vertices == 0:
        # Point group Ci (-1) has no vertices
        rx = Vector3d.xvector() - center
    else:
        rx = Vector3d.zvector() - center  # North pole to sector center
    azimuth = _calculate_azimuth(center, rx, v)
    if n_vertices != 0:
        azimuth = _correct_azimuth(azimuth, sector, rx)

    # Polar coordinate
    if np.count_nonzero(sector.dot(center)) == 0:
        polar = center.angle_with(v) / np.pi
    else:
        # Normal to plane containing sector center and crystal direction
        v_center_normal = v.cross(center).unit
        polar = np.full(v.shape, np.inf)
        for normal in sector:
            boundary_points = v_center_normal.cross(normal).unit
            # Some boundary points are zero vectors
            with np.errstate(invalid="ignore"):
                distances_polar = (-v).angle_with(boundary_points)
                distances_polar /= (-center).angle_with(boundary_points)
            distances_polar[np.isnan(distances_polar)] = 1
            polar = np.minimum(polar, distances_polar)  # Element-wise

    return azimuth, polar


def rgb_from_polar_coordinates(azimuth: np.ndarray, polar: np.ndarray) -> np.ndarray:
    """Calculate RGB colors from polar coordinates.

    Parameters
    ----------
    azimuth
        Azimuthal coordinate(s).
    polar
        Polar coordinate(s).

    Returns
    -------
    rgb
        Color(s).
    """
    angle = np.mod(azimuth / (2 * np.pi), 1)
    h, s, v = hsl_to_hsv(angle, 1, polar)
    return mcolors.hsv_to_rgb(np.stack((h, s, v), axis=-1))


def hsl_to_hsv(
    hue: np.ndarray | float,
    saturation: np.ndarray | float,
    lightness: np.ndarray | float,
) -> tuple[np.ndarray | float, np.ndarray | float, np.ndarray | float]:
    """Convert color described by HSL (hue, saturation and lightness) to
    HSV (hue, saturation and value).

    Adapted from MTEX' function :code:`hsl2hsv`.

    Parameters
    ----------
    hue
        Hue(s). Not changed by the function, but included in input and
        output for convenience.
    saturation
        Saturation value(s).
    lightness
        Lightness value(s).

    Returns
    -------
    hue
        The same copy of hue(s) as input.
    saturation2
        Adjusted saturation value(s).
    value
        Value(s).
    """
    l2 = 2 * lightness
    s2 = saturation * np.where(l2 <= 1, l2, 2 - l2)
    saturation2 = (2 * s2) / (l2 + s2)
    saturation2[np.isnan(saturation2)] = 0
    value = (l2 + s2) / 2
    return hue, saturation2, value


def _calculate_azimuth(center: Vector3d, rx: Vector3d, v: Vector3d) -> np.ndarray:
    """Calculate the azimuthal coordinate of the polar coordinates of a
    fundamental sector (inverse pole figure).

    Taken from MTEX' :code:`calcAngle` function.

    Parameters
    ----------
    center
    rx
    v

    Returns
    -------
    azimuth
    """
    rx = (rx - rx.dot(center) * center).unit  # Orthogonal to center
    ry = center.cross(rx).unit  # Perpendicular to rx
    distances_azimuthal = (v - center).unit
    azimuth = np.arctan2(ry.dot(distances_azimuthal), rx.dot(distances_azimuthal))
    azimuth = np.mod(azimuth, 2 * np.pi)
    azimuth[np.isnan(azimuth)] = 0
    return azimuth


def _correct_azimuth(
    azimuth: np.ndarray, sector: FundamentalSector, rx: Vector3d
) -> np.ndarray:
    """Correct the azimuthal coordinate of the polar coordinates of a
    fundamental sector (inverse pole figure) when vertices are not
    an equal distance from each other.

    This ensures for example that the entirely blue part of the m-3m
    sector is in the [111] direction, and not in the direction
    corresponding to the sector position perpendicular to the
    [001]-[101] line (straight north from the barycenter).

    Taken from MTEX' :code:`correctAngle` function.

    Parameters
    ----------
    azimuth
    sector
    rx

    Returns
    -------
    azimuth
        Corrected azimuthal coordinates.
    """
    center = sector.center.unit
    vertices = sector.vertices.unit

    m = 1000
    azimuth2 = np.linspace(0, 2 * np.pi, m)
    rot = Rotation.from_axes_angles(center, azimuth2[:-1])
    normals = rot * rx.cross(center).unit

    # Distance between center and all boundary points
    cross_outer = Vector3d.zero((sector.size, m - 1))
    for i in range(sector.size):
        cross_outer[i] = sector[i].cross(normals)
    polar = np.min(cross_outer.angle_with(center), axis=0)

    if vertices.size == 3:
        angle = _calculate_azimuth(center, rx, vertices)
        azimuth3 = np.round(m * np.sort(angle) / (2 * np.pi)).astype(int)
        if azimuth3[0] < 10:
            azimuth3 = np.delete(azimuth3, 0)
        idx = np.arange(azimuth3[0])
        polar[idx] /= np.sum(polar[idx]) / 3
        idx = np.arange(azimuth3[0], azimuth3[1])
        polar[idx] /= np.sum(polar[idx]) / 3
        idx = np.arange(azimuth3[1], polar.size)
        polar[idx] /= np.sum(polar[idx]) / 3

    polar = 2 * np.pi * np.cumsum(np.append(np.zeros(1), polar / np.sum(polar)))
    azimuth = np.interp(azimuth, azimuth2, polar)

    return azimuth
