import sys
import functools
import numpy as np
import cv2
import cmapy
from PIL import Image
import matplotlib



# BUGFIX in cmapy.py
def cmap(cmap_name, rgb_order=False):
    """
    Extract colormap color information as a LUT compatible with cv2.applyColormap().
    Default channel order is BGR.

    Args:
        cmap_name: string, name of the colormap.
        rgb_order: boolean, if false or not set, the returned array will be in
                   BGR order (standard OpenCV format). If true, the order
                   will be RGB.

    Returns:
        A numpy array of type uint8 containing the colormap.
    """

    c_map = matplotlib.colormaps.get_cmap(cmap_name)
    rgba_data = matplotlib.cm.ScalarMappable(cmap=c_map).to_rgba(
        np.arange(0, 1.0, 1.0 / 256.0), bytes=True
    )
    rgba_data = rgba_data[:, 0:-1].reshape((256, 1, 3))

    # Convert to BGR (or RGB), uint8, for OpenCV.
    cmap = np.zeros((256, 1, 3), np.uint8)

    if not rgb_order:
        cmap[:, :, :] = rgba_data[:, :, ::-1]
    else:
        cmap[:, :, :] = rgba_data[:, :, :]

    return cmap

# If python 3, redefine cmap() to use lru_cache.
if sys.version_info > (3, 0):
    cmap = functools.lru_cache(maxsize=200)(cmap)



def alpha_composite(img, msk, opacity=0.5, colormap=None, alpha_image=None, alpha_mask=None, red_mask=False):
    """Alpha composite an RGBA image (img) and a grayscale mask (msk).
    - If alpha_image is None, img's alpha channel is used (or, if not present,
    initialized to all 255).
    - If alpha_mask is None, msk is overlaid on img only where img's alpha
    channel is not 0.
    - If alpha_mask is not None, the above behavior is overridden and msk is
    overlaid on img only where alpha_mask is not 0."""
    # only HWC numpy arrays allowed
    assert isinstance(img, np.ndarray), f'Input image must be a numpy array. Got {type(img)}'
    assert isinstance(msk, np.ndarray), f'Input mask must be a numpy array. Got {type(msk)}'
    if alpha_mask is not None:
        assert isinstance(alpha_mask, np.ndarray), f'Alpha mask must be a numpy array. Got {type(alpha_mask)}'
        assert alpha_mask.dtype in [np.float32, bool], f'Alpha mask must be of type np.float32 or bool. Got {alpha_mask.dtype}'
        assert alpha_mask.shape[2] == 1, f'Alpha mask must be formatted as HWC, with C = 1. Got a shape of {msk.shape}'
    assert img.shape[2] in [3,4], f'Input image must be formatted as HWC, with C = 3,4. Got a shape of {img.shape}'
    assert msk.shape[2] == 1, f'Input mask must be formatted as HWC, with C = 1. Got a shape of {msk.shape}'
    assert (opacity >= 0) and (opacity <= 1), f'Mask opacity must be between 0 and 1. Got {opacity}'
    
    # to avoid modifying the original arrays
    img = img.copy()
    msk = msk.copy()

    if img.shape[2] == 3:
        # add alpha channel to img
        img = np.concatenate([
            img,
            np.full((img.shape[0], img.shape[1], 1), 255, dtype=np.uint8)
            ], axis=-1)

    if alpha_image is None:
        # initialize alpha_image to all Trues
        alpha_image = img[:,:,[3]]
    # convert alpha image to bool
    alpha_image = alpha_image.astype(bool)

    if alpha_mask is None:
        # initialize alpha_mask to alpha_image
        alpha_mask = alpha_image    # so that alpha_mask is AT LEAST as restrictive as alpha_image
    # convert alpha mask to bool
    alpha_mask = alpha_mask.astype(bool)


    if msk.dtype != np.uint8:
        # convert mask to a uint8 grayscale image ([0,1] -> [0,255])
        # NB: normalize the pixels of the mask we are interested in to [0,1]
        # before passing it as input!!!
        msk = (msk * 255).astype(np.uint8)

    # convert mask from grayscale to RGBA
    msk = cv2.cvtColor(msk, cv2.COLOR_GRAY2RGBA)

    if colormap is not None:
        # apply specified colormap to msk
        # NB: values near 0 will be converted to the first colors of the chosen
        # colormap, whereas values near 255 will be converted to the last colors
        msk[:,:,:3] = cmapy.colorize(msk[:,:,:3], colormap, rgb_order=True)
    elif red_mask:
        # convert white to red
        msk[:,:,[1,2]] = 0


    # apply alpha_image to img's alpha channel
    img[:,:,[3]] = (alpha_image * img[:,:,[3]]).astype(np.uint8)

    # apply alpha_mask and opacity to msk's alpha channel
    msk[:,:,[3]] = (alpha_mask * opacity * msk[:,:,[3]]).astype(np.uint8)

    # alpha compositing
    img_pil = Image.fromarray(img)
    msk_pil = Image.fromarray(msk)
    img_pil.alpha_composite(msk_pil)

    return np.array(img_pil)