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)