__author__ = 'maartenbreddels'
import PIL.Image
import PIL.ImageDraw
try:
    from StringIO import StringIO
    py3 = False
except ImportError:
    from io import StringIO, BytesIO
    py3 = True
import collections
import numpy as np
import matplotlib.colors
import warnings
from base64 import b64encode


def rgba_2_pil(rgba):
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        im = PIL.Image.fromarray(rgba[::-1], "RGBA")  # , "RGBA", 0, -1)
    return im


def pil_2_data(im, format="png"):
    if py3:  # py3 case
        f = BytesIO()
    else:
        f = StringIO()
    im.save(f, format)
    return f.getvalue()


def rgba_to_url(rgba):
    bit8 = rgba.dtype == np.uint8
    if not bit8:
        rgba = (rgba * 255.).astype(np.uint8)
    im = rgba_2_pil(rgba)
    data = pil_2_data(im)
    data = b64encode(data)
    if py3:
        data = data.decode("ascii")
    imgurl = "data:image/png;base64," + data + ""
    return imgurl


pdf_modes = collections.OrderedDict()
pdf_modes["multiply"] = np.multiply
pdf_modes["screen"] = lambda a, b: a + b - a * b
pdf_modes["darken"] = np.minimum
pdf_modes["lighten"] = np.maximum

cairo_modes = collections.OrderedDict()
cairo_modes["saturate"] = (lambda aA, aB: np.minimum(1, aA + aB),
                           lambda aA, xA, aB, xB, aR: (np.minimum(aA, 1 - aB) * xA + aB * xB) / aR)
cairo_modes["add"] = (lambda aA, aB: np.minimum(1, aA + aB),
                      lambda aA, xA, aB, xB, aR: (aA * xA + aB * xB) / aR)

modes = list(pdf_modes.keys()) + list(cairo_modes.keys())


def background(shape, color="white", alpha=1, bit8=True):
    rgba = np.zeros(shape + (4,))
    rgba[:] = np.array(matplotlib.colors.colorConverter.to_rgba(color))
    rgba[..., 3] = alpha
    if bit8:
        return (rgba * 255).astype(np.uint8)
    else:
        return rgba


def fade(image_list, opacity=0.5, blend_mode="multiply"):
    result = image_list[0] * 1.
    for i in range(1, len(image_list)):
        result[result[..., 3] > 0, 3] = 1
        layer = image_list[i] * 1.0
        layer[layer[..., 3] > 0, 3] = opacity
        result = blend([result, layer], blend_mode=blend_mode)
    return result


def blend(image_list, blend_mode="multiply"):
    bit8 = image_list[0].dtype == np.uint8
    image_list = image_list[::-1]
    rgba_dest = image_list[0] * 1
    if bit8:
        rgba_dest = (rgba_dest / 255.).astype(np.float)
    for i in range(1, len(image_list)):
        rgba_source = image_list[i]
        if bit8:
            rgba_source = (rgba_source / 255.).astype(np.float)
        # assert rgba_source.dtype == image_list[0].dtype, "images have different types: first has %r, %d has %r" % (image_list[0].dtype, i, rgba_source.dtype)

        aA = rgba_source[:, :, 3]
        aB = rgba_dest[:, :, 3]

        if blend_mode in pdf_modes:
            aR = aA + aB * (1 - aA)
        else:
            aR = cairo_modes[blend_mode][0](aA, aB)
        mask = aR > 0
        for c in range(3):  # for r, g and b
            xA = rgba_source[..., c]
            xB = rgba_dest[..., c]
            if blend_mode in pdf_modes:
                f = pdf_modes[blend_mode](xB, xA)
                with np.errstate(divide='ignore', invalid='ignore'):  # these are fine, we are ok with nan's in vaex
                    result = ((1. - aB) * aA * xA + (1. - aA) * aB * xB + aA * aB * f) / aR
            else:
                result = cairo_modes[blend_mode][1](aA, xA, aB, xB, aR)
            with np.errstate(divide='ignore', invalid='ignore'):  # these are fine, we are ok with nan's in vaex
                result = (np.minimum(aA, 1 - aB) * xA + aB * xB) / aR
            # print(result)
            rgba_dest[:, :, c][(mask, )] = np.clip(result[(mask,)], 0, 1)
        rgba_dest[:, :, 3] = np.clip(aR, 0., 1)
    rgba = rgba_dest
    if bit8:
        rgba = (rgba * 255).astype(np.uint8)
    return rgba


def monochrome(I, color, vmin=None, vmax=None):
    """Turns a intensity array to a monochrome 'image' by replacing each intensity by a scaled 'color'

    Values in I between vmin  and vmax get scaled between 0 and 1, and values outside this range are clipped to this.

    Example
    >>> I = np.arange(16.).reshape(4,4)
    >>> color = (0, 0, 1) # red
    >>> rgb = vx.image.monochrome(I, color) # shape is (4,4,3)

    :param I: ndarray of any shape (2d for image)
    :param color: sequence of a (r, g and b) value
    :param vmin: normalization minimum for I, or np.nanmin(I) when None
    :param vmax: normalization maximum for I, or np.nanmax(I) when None
    :return:
    """
    if vmin is None:
        vmin = np.nanmin(I)
    if vmax is None:
        vmax = np.nanmax(I)
    normalized = (I - vmin) / (vmax - vmin)
    return np.clip(normalized[..., np.newaxis], 0, 1) * np.array(color)


def polychrome(I, colors, vmin=None, vmax=None, axis=-1):
    """Similar to monochrome, but now do it for multiple colors

    Example
    >>> I = np.arange(32.).reshape(4,4,2)
    >>> colors = [(0, 0, 1), (0, 1, 0)] # red and green
    >>> rgb = vx.image.polychrome(I, colors) # shape is (4,4,3)

    :param I: ndarray of any shape (3d will result in a 2d image)
    :param colors: sequence of [(r,g,b), ...] values
    :param vmin: normalization minimum for I, or np.nanmin(I) when None
    :param vmax: normalization maximum for I, or np.nanmax(I) when None
    :param axis: axis which to sum over, by default the last
    :return:
    """
    axes_length = len(I.shape)
    allaxes = list(range(axes_length))
    otheraxes = list(allaxes)
    otheraxes.remove((axis + axes_length) % axes_length)
    otheraxes = tuple(otheraxes)

    if vmin is None:
        vmin = np.nanmin(I, axis=otheraxes)
    if vmax is None:
        vmax = np.nanmax(I, axis=otheraxes)
    normalized = (I - vmin) / (vmax - vmin)
    return np.clip(normalized, 0, 1).dot(colors)

# c_r + c_g + c_b
