from __future__ import annotations

import copy
import sys
from typing import Callable

import matplotlib.colors as pltc
import matplotlib.pyplot as plt
import numpy as np

from plotpy.mathutils.colormap import (
    DEFAULT_COLORMAPS,
    DEFAULT_COLORMAPS_PATH,
    CmapDictType,
    EditableColormap,
    save_colormaps,
)

PERCEPTUALLY_UNIFORM_CMAPS = ["viridis", "plasma", "inferno", "magma", "cividis"]
SEQUENTIAL_CMAPS = [
    "Greys",
    "Purples",
    "Blues",
    "Greens",
    "Oranges",
    "Reds",
    "YlOrBr",
    "YlOrRd",
    "OrRd",
    "PuRd",
    "RdPu",
    "BuPu",
    "GnBu",
    "PuBu",
    "YlGnBu",
    "PuBuGn",
    "BuGn",
    "YlGn",
]
SEQUENTIAL_CMAPS2 = [
    "binary",
    "gist_yarg",
    "gist_gray",
    "gray",
    "bone",
    "pink",
    "spring",
    "summer",
    "autumn",
    "winter",
    "cool",
    "Wistia",
    "hot",
    "afmhot",
    "gist_heat",
    "copper",
]
DIVERGING_CMAPS = [
    "PiYG",
    "PRGn",
    "BrBG",
    "PuOr",
    "RdGy",
    "RdBu",
    "RdYlBu",
    "RdYlGn",
    "Spectral",
    "coolwarm",
    "bwr",
    "seismic",
]
CYCLIC_CMAPS = ["twilight", "twilight_shifted", "hsv"]
QUALITATIVE_CMAPS = [
    "Pastel1",
    "Pastel2",
    "Paired",
    "Accent",
    "Dark2",
    "Set1",
    "Set2",
    "Set3",
    "tab10",
    "tab20",
    "tab20b",
    "tab20c",
]
MISCELLANEOUS_CMAPS = [
    "flag",
    "prism",
    "ocean",
    "gist_earth",
    "terrain",
    "gist_stern",
    "gnuplot",
    "gnuplot2",
    "CMRmap",
    "cubehelix",
    "brg",
    "gist_rainbow",
    "rainbow",
    "jet",
    "turbo",
    "nipy_spectral",
    "gist_ncar",
]

SORTED_MATPLOTLIB_COLORMAPS: list[str] = [
    *PERCEPTUALLY_UNIFORM_CMAPS,
    *SEQUENTIAL_CMAPS,
    *SEQUENTIAL_CMAPS2,
    *DIVERGING_CMAPS,
    *CYCLIC_CMAPS,
    *QUALITATIVE_CMAPS,
    *MISCELLANEOUS_CMAPS,
]


def rgb_colors_to_hex_list(
    colors: list[tuple[int, int, int]],
) -> list[tuple[float, str]]:
    """Convert a list of RGB colors to a list of tuples with the position of the color
    and the color in hex format. Positions evenly distributed between 0 and 1.

    Args:
        colors: list of RGB colors

    Returns:
        list of tuples with the position of the color and the color in hex format
    """
    return [(i / len(colors), pltc.to_hex(color)) for i, color in enumerate(colors)]


def _interpolate(
    val: float, vmin: tuple[float, float, float], vmax: tuple[float, float, float]
):
    """Interpolate between two level of a color.

    Args:
        val: value to interpolate
        vmin: R, G or B tuple from a matplotlib segmented colormap
        vmax: R, G or B tuple from matplotlib segmented colormap

    Returns:
        The interpolated R, G or B component
    """
    interp = (val - vmin[0]) / (vmax[0] - vmin[0])
    return (1 - interp) * vmin[1] + interp * vmax[2]


def std_segmented_cmap_to_hex_list(cmdata: dict[str, list[tuple[float, float, float]]]):
    """Convert a matplotlib segmented colormap to a list of tuples with the position of
    the color and the color in hex format.

    Args:
        cmdata: segmented colormap data

    Returns:
        list of tuples with the position of the color and the color in hex format
    """
    colors: list[tuple[float, str]] = []
    red = np.array(cmdata["red"])
    green = np.array(cmdata["green"])
    blue = np.array(cmdata["blue"])
    indices = sorted(set(red[:, 0]) | set(green[:, 0]) | set(blue[:, 0]))
    for i in indices:
        idxr = red[:, 0].searchsorted(i)
        idxg = green[:, 0].searchsorted(i)
        idxb = blue[:, 0].searchsorted(i)
        compr = _interpolate(i, red[idxr - 1], red[idxr])
        compg = _interpolate(i, green[idxg - 1], green[idxg])
        compb = _interpolate(i, blue[idxb - 1], blue[idxb])
        colors.append((i, pltc.to_hex((compr, compg, compb))))
    return colors


InterpFuncT = Callable[[np.ndarray], np.ndarray]


def func_segmented_cmap_to_hex_list(
    n: int,
    cmap: pltc.LinearSegmentedColormap,
) -> list[tuple[float, str]]:
    """Convert a matplotlib segmented colormap to a list of tuples with the position of
    the color and the color in hex format. The input colormap contains function for each
    color RGB component instead of a list of colors.

    Args:
        n: number of colors to generate
        cmap: segmented colormap

    Returns:
        list of tuples with the position of the color and the color in hex format
    """
    colors = []
    arr = np.linspace(0, 1, n, dtype=float)
    colors = [(i, pltc.to_hex(rgba)) for i, rgba in zip(arr, cmap(arr))]
    return colors


def continuous_to_descrete_cmap(cmap: EditableColormap) -> EditableColormap:
    """Convert a continuous colormap to a descrete one.

    Args:
        cmap: colormap to convert

    Returns:
        descrete colormap
    """
    raw_cmap: tuple[tuple[float, str], ...] = cmap.to_tuples()
    new_raw_cmap: list[tuple[float, str]] = [raw_cmap[0]]
    n = len(raw_cmap)
    coeff = (n - 1) / n
    for i, (pos, color) in enumerate(raw_cmap[1:]):
        prev_pos, prev_color = raw_cmap[i]
        curr_pos, curr_color = pos, color
        new_pos = curr_pos * coeff
        new_raw_cmap.append((new_pos, prev_color))
        new_raw_cmap.append((new_pos, curr_color))
    new_raw_cmap.append(raw_cmap[-1])

    return EditableColormap.from_iterable(new_raw_cmap, name=cmap.name)


def sort_mpl_colormaps(colormaps: CmapDictType) -> CmapDictType:
    """Filter and sort input colormaps to follow the same order (by category) as in the
     matplotlib colormaps documentation. Colormaps not found in the matplotlib
     are filtered out.

    Args:
        colormaps: Dictionnary of colormaps to extract and order

    Returns:
        Filtered and sorted colormaps dictionnary
    """
    ordered_colormaps: CmapDictType = {}
    lower_cmap_names = [cm.lower() for cm in SORTED_MATPLOTLIB_COLORMAPS]
    for lower_name in lower_cmap_names:
        cmap = colormaps.get(lower_name, None)
        if lower_name.endswith("_r"):
            continue
        if cmap is None:
            print(f"Colormap {lower_name} not found in input colormaps.")
            continue
        ordered_colormaps[lower_name] = cmap
    return ordered_colormaps


def append_non_mpl_colormaps(mpl_colormaps: CmapDictType, colormaps: CmapDictType):
    """Append colormaps not found in the matplotlib colormaps to the input colormaps.
     Mutate the input in place.

    Args:
        mpl_colormaps: dictionnary of matplotlib colormaps. Mutated in place.
        colormaps: dictionnary of colormaps to append to the matplotlib colormaps
    """
    colormap_names = set(SORTED_MATPLOTLIB_COLORMAPS)
    for colormap in colormaps.values():
        if colormap.name not in colormap_names:
            print(f"{colormap} not in matplotlib colormaps.")
            mpl_colormaps[colormap.name.lower()] = colormap


def main(cmaps: CmapDictType, out_json_path: str = DEFAULT_COLORMAPS_PATH):
    new_cmaps: dict[str, list[tuple[float, str]]] = {}

    # Uniform colormaps with a .colors attribute that return a list of RGB colors
    cmaps_with_colors = [
        "magma",
        "viridis",
        "inferno",
        "plasma",
        "cividis",
    ]

    # Discrete colormaps, same as uniform colormaps but the colormap must be post
    # processed to become descrete
    descrete_cmaps = [
        "Pastel1",
        "Pastel2",
        "Paired",
        "Accent",
        "Dark2",
        "Set1",
        "Set2",
        "Set3",
    ]

    cmaps_with_colors.extend(descrete_cmaps)

    # Colormaps with a _segmented_data attribute that contains the R, G and B components
    # as lists of tuples
    segmented_cmaps = [
        "coolwarm",
        "bwr",
        "seismic",
    ]

    # Colormaps with a _segmentdata attribute that contains the R, G and B components as
    # functions that return the color for a given position
    interp_cmaps = ["gnuplot2", "CMRmap", "rainbow", "turbo", "afmhot"]

    for cm_name in cmaps_with_colors:
        cmap = plt.get_cmap(cm_name)
        new_cmaps[cm_name] = rgb_colors_to_hex_list(cmap.colors)

    for cm_name in descrete_cmaps:
        cmap = EditableColormap.from_iterable(new_cmaps[cm_name], name=cm_name)
        new_cmaps[cm_name] = list(continuous_to_descrete_cmap(cmap).to_tuples())

    for cm_name in segmented_cmaps:
        cmap = plt.get_cmap(cm_name)
        new_cmaps[cm_name] = std_segmented_cmap_to_hex_list(cmap._segmentdata)

    n = 128

    for cm_name in interp_cmaps:
        cmap = plt.get_cmap(cm_name)
        new_cmaps[cm_name] = func_segmented_cmap_to_hex_list(n, cmap)

    for name, raw_cm in new_cmaps.items():
        cmaps[name.lower()] = EditableColormap.from_iterable(raw_cm, name=name)

    ordered_cmaps = sort_mpl_colormaps(cmaps)
    append_non_mpl_colormaps(ordered_cmaps, cmaps)
    save_colormaps(out_json_path, ordered_cmaps)


if __name__ == "__main__":
    cmaps = copy.deepcopy(DEFAULT_COLORMAPS)
    out_json = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_COLORMAPS_PATH
    main(cmaps, out_json)