"""Module containing the palette definitions flyr supports and the function
with which to use it.

Specifically see the `map_colors` function, which uses any of the palettes
defined in the `palettes` dictionary.
"""

from typing import Dict, List, Literal, Tuple, Union

import numpy as np

import flyr.normalization as norm
from flyr.dtypes import Generic2DArray, RGBArray
from flyr.palettes.cividis import cividis
from flyr.palettes.copper import copper
from flyr.palettes.gist_earth import gist_earth
from flyr.palettes.gist_rainbow import gist_rainbow
from flyr.palettes.hot import hot
from flyr.palettes.inferno import inferno
from flyr.palettes.jet import jet
from flyr.palettes.magma import magma
from flyr.palettes.ocean import ocean
from flyr.palettes.plasma import plasma
from flyr.palettes.rainbow import rainbow
from flyr.palettes.terrain import terrain
from flyr.palettes.turbo import turbo
from flyr.palettes.viridis import viridis

PaletteBuiltIn = Union[
    Literal["cividis"],
    Literal["copper"],
    Literal["gist_earth"],
    Literal["gist_rainbow"],
    Literal["hot"],
    Literal["inferno"],
    Literal["jet"],
    Literal["magma"],
    Literal["ocean"],
    Literal["plasma"],
    Literal["rainbow"],
    Literal["terrain"],
    Literal["turbo"],
    Literal["viridis"],
]
Palette = Union[
    PaletteBuiltIn,
    Literal["embedded"],
    Literal["grayscale"],
    Literal["grayscale-inverted"],
]

PaletteColors = List[Tuple[int, int, int]]

palettes: Dict[PaletteBuiltIn, List[Tuple[int, int, int]]] = {
    "cividis": cividis,
    "copper": copper,
    "gist_earth": gist_earth,
    "gist_rainbow": gist_rainbow,
    "hot": hot,
    "inferno": inferno,
    "jet": jet,
    "magma": magma,
    "ocean": ocean,
    "plasma": plasma,
    "rainbow": rainbow,
    "terrain": terrain,
    "turbo": turbo,
    "viridis": viridis,
}


def map_colors(
    thermal: Generic2DArray[np.float64],
    palette: Union[PaletteBuiltIn, PaletteColors] = "jet",
    histogram_equalized: bool = False,
) -> RGBArray:
    """Colors a normalized image array.

    Parameters
    ----------
    thermal: Generic2DArray[np.float64]
        A normalized array, that is all values are between 0 and 1.

    Returns
    -------
    RGBArray
        A 3D array with channels as its last dimension, containing integers
        between 0 and 255.
    """
    if isinstance(palette, list):
        colors = palette
    elif isinstance(palette, str):
        colors = palettes[palette]

    colored = np.zeros(thermal.shape + (3,), dtype=np.uint8)
    num_colors = len(colors)
    if histogram_equalized:
        thermal = norm.equalize_histogram(thermal, num_colors)

    for idx, color in enumerate(colors):
        lower = idx / num_colors
        upper = (idx + 1) / num_colors
        colored[(lower <= thermal) & (thermal <= upper)] = color

    return colored
