# Copyright (c) 2022 The xlandsat developers.
# Distributed under the terms of the MIT License.
# SPDX-License-Identifier: MIT
"""
Operations to enhance composites. Basically wrappers for skimage.exposure.
"""
import numpy as np
import skimage.color
import skimage.exposure


def equalize_histogram(composite, kernel_size=None, clip_limit=0.01):
    """
    Adaptive histogram equalization for a composite

    Use this function to enhance the contrast of a composite when there are a
    few very dark or very light patches that dominate the color range. Use this
    instead of rescaling intensity (contrast stretching) to try to preserve
    some detail in the light/dark patches.

    If the composite has an alpha channel (transparency), it will be copied to
    the output intact.

    .. warning::

        Results can be very bad if there are missing values (NaNs) in the
        composite. Use :func:`xlandsat.interpolate_missing` on the scene first
        (before creating the composite) if that is the case.

    Parameters
    ----------
    composite : :class:`xarray.DataArray`
        A composite, as generated by :func:`xlandsat.composite`.
    kernel_size : int or None
        The size of region used in the algorithm. See
        :func:`skimage.exposure.equalize_adapthist` for details and defaults.
    clip_limit : float
        Clipping limit, normalized between 0 and 1 (higher values give more
        contrast).

    Returns
    -------
    equalized_composite : :class:`xarray.DataArray`
        The composite after equalization, scaled back to the range of the
        original composite.

    Notes
    -----

    This function first converts the composite from the RGB color space to the
    `HSV color space <https://en.wikipedia.org/wiki/HSL_and_HSV>`__. Then, it
    applies :func:`skimage.exposure.equalize_adapthist` to the values
    (intensity) channel. Finally, the composite is converted back into the RGB
    color space.
    """
    result = composite.copy(deep=True)
    # Make sure the input is in the 0-255 range for the color space transform
    hsv = skimage.color.rgb2hsv(
        skimage.exposure.rescale_intensity(result.values[:, :, :3], out_range="uint8")
    )
    hsv[:, :, 2] = skimage.exposure.equalize_adapthist(
        hsv[:, :, 2],
        kernel_size=kernel_size,
        clip_limit=clip_limit,
    )
    result.values[:, :, :3] = skimage.exposure.rescale_intensity(
        skimage.color.hsv2rgb(hsv),
        out_range=str(composite.values.dtype),
    )
    return result


def adjust_l1_colors(composite, percentile=0.1):
    """
    Adjust the colors in an RGB composite from Level 1 data

    Corrects the balance of the red, green, and blue bands in an RGB (true
    color) composite made from Level 1 data (without atmospheric correction).
    This is **not as accurate as atmospheric correction** but can lead to nicer
    images in places where the atmospheric correction causes artifacts.

    **Do not use the output for calculating indices.**

    Parameters
    ----------
    composite : :class:`xarray.DataArray`
        A composite, as generated by :func:`xlandsat.composite`.
    percentile : float
        The percentile range to use for the intensity rescaling. Will use the
        ``percentile`` as the lower bound and ``100 - percentile`` for the
        upper bound. Default is 0.1%.

    Returns
    -------
    adjusted_composite : :class:`xarray.DataArray`
        The composite after color adjustment, scaled back to the range of the
        original composite.

    Notes
    -----

    The correction raises each channel to the power of 1/2.2 (inverse of the
    sRGB gamma function) and then rescales the transformed intensity to the
    given percentile range.
    """
    output = composite.copy()
    for i, channel in enumerate(["red", "green", "blue"]):
        scaled = composite.sel(channel=channel).values ** (1 / 2.2)
        vmin, vmax = np.percentile(scaled, (percentile, 100 - percentile))
        output.values[:, :, i] = skimage.exposure.rescale_intensity(
            scaled,
            in_range=(vmin, vmax),
            out_range=str(composite.values.dtype),
        )
    return output
