from typing import Literal, Union

import numpy as np

from flyr.dtypes import Generic1DArray, Generic2DArray

NormalizationMode = Union[Literal["minmax"], Literal["percentiles"], Literal["hist-eq"]]


def by_minmax(
    min_t: float, max_t: float, thermal: Generic2DArray[np.float64]
) -> Generic2DArray[np.float64]:
    """Normalize the given array to be between 0 and 1 using the given minimum
    and maximum values.

    Parameters
    ----------
    min_t: float
        The value to clip the input array to and which should map to 0.
    max_t: float
        The value to clip the input array to and which should map to 1.
    thermal: Generic2DArray[np.float64]
        The array to normalize.

    Returns
    -------
    Generic2DArray[np.float64]
        An array with all values between 0 and 1
    """
    err_msg = f"Minimum value {min_t} should be smaller than maximum value {max_t}"
    assert min_t < max_t, err_msg
    thermal = np.clip(thermal, min_t, max_t)
    thermal = (thermal - min_t) / (max_t - min_t)
    return thermal


def by_histogram_equalization(
    min_t: float,
    max_t: float,
    thermal: Generic2DArray[np.float64],
    num_colors: int = 256,
) -> Generic2DArray[np.float64]:
    """Normalize the given array using histogram equalization, improving
    contrast.

    Warning: Improves contrast, but changes the apparant temperature.

    Parameters
    ----------
    min_t: float
        The value to clip the input array to and which should map to 0.
    max_t: float
        The value to clip the input array to and which should map to 1.
    thermal: Generic2DArray[np.float64]
        The array to normalize.
    num_colors: int
        How many colors to equalize the histogram over. In other words, the
        number of histogram bins.

    Returns
    -------
    Generic2DArray[np.float64]
        An array with all values between 0 and 1
    """
    err_msg = f"Minimum value {min_t} should be smaller than maximum value {max_t}"
    assert min_t < max_t, err_msg

    normalized = by_minmax(min_t, max_t, thermal)
    return equalize_histogram(normalized, num_colors)


def equalize_histogram(normalized: Generic2DArray, num_colors: int):
    max_color_value = num_colors - 1
    color_chances, colored = __calc_color_chances(normalized, num_colors)
    new_colors = __calc_new_colors(color_chances)
    return new_colors[colored] / max_color_value


def __calc_color_chances(
    normalized: Generic2DArray[np.float64], num_colors: int
) -> tuple[Generic1DArray[np.float64], Generic2DArray[np.int64]]:
    """

    Parameters
    ----------
    normalized
        A 2D float numpy array of the thermal data normalized to between 0 and 1.
    num_colors: int
        The number of colors we want to calculate the occurrance chance of.

    Returns
    -------
    tuple[color chances, color map]
        A 2-tuple, where the first item is an array whose value indicates
        the chance to occur of the color identified by the index in this array.
        The second item is a map, derived from the normalized array, where
        the values correspond to color indices from the color chances array.
    """
    max_color = num_colors - 1
    color_map = (normalized * max_color).astype(int)

    [histogram, _] = np.histogram(color_map, bins=num_colors, range=(0, max_color))
    color_chances = histogram / normalized.size

    return color_chances, color_map


def __calc_new_colors(
    color_chances: Generic1DArray[np.float64],
) -> Generic2DArray[np.int64]:
    """Calculate a new color based on the current color.

    s = T(i) where i is a pixel's intensity
    0 <= i <= 255
    T(i) <= T(i+1)
    0 <= T(i) <= 255
    """
    max_color = color_chances.shape[0] - 1
    cdf = np.cumsum(color_chances)
    return np.floor(max_color * cdf).astype(int)


def by_percentiles(
    min_p: float, max_p: float, thermal: Generic2DArray[np.float64]
) -> Generic2DArray[np.float64]:
    """Normalizes the given thermal array using the values at the given minimum
    and maximum percentiles.

    Parameters
    ----------
    min_p: float
        Value between 0 and 1 that is lower than `max_p`. Indicates the
        percentile from which to take the minimum normalization value.
    max_p: float
        Value between 0 and 1 that is greater than `min_p`. Indicates the
        percentile from which to take the maximum normalization value.
    thermal: Generic2DArray[np.float64]
        Array to normalize.

    Returns
    -------
    Generic2DArray[np.float64]
        The normalized array, all np.float64 values between 0 and 1.
    """
    assert 0.0 <= min_p < max_p <= 1.0
    return by_minmax(
        np.percentile(thermal, int(min_p * 100)).item(),
        np.percentile(thermal, int(max_p * 100)).item(),
        thermal,
    )
