import glob
import numpy as np
import os
import pylab
from OpenImageIO import FLOAT, ImageInput
from scipy.spatial import Delaunay

import colour
from colour.plotting import *
from colour.utilities import message_box

figure_dimensions = 28, 14
pylab.rcParams['figure.figsize'] = figure_dimensions


def exr_image_as_array(path):
    image = ImageInput.open(path)
    specification = image.spec()

    return np.array(image.read_image(FLOAT)).reshape((specification.height,
                                                      specification.width,
                                                      specification.nchannels))

@figure_size(figure_dimensions)
def image_plot(image,
               transfer_function=colour.sRGB_COLOURSPACE.transfer_function):
    vectorised_oecf = np.vectorize(transfer_function)
    image = np.clip(vectorised_oecf(image), 0, 1)
    pylab.imshow(image)

    settings = {'no_ticks': True,
                'bounding_box': [0, 1, 0, 1],
                'bbox_inches': 'tight',
                'pad_inches': 0}

    aspect(**settings)
    display(**settings)

directory = 'color_checker_comparisons_v0002'
exr_images = glob.glob('{0}/*medium.exr'.format(directory))

colourspaces_mapping = {'p3_medium.exr': colour.DCI_P3_COLOURSPACE,
                        'rec709_medium.exr': colour.REC_709_COLOURSPACE}


def RGB_to_xy(RGB, colourspace):
    return colour.XYZ_to_xy(
                colour.RGB_to_XYZ(RGB,
                                  colourspace.whitepoint,
                                  colourspace.whitepoint,
                                  colourspace.RGB_to_XYZ_matrix))


def chromaticities_plot(xy, colourspace):
    colourspaces_CIE_1931_chromaticity_diagram_plot(
        ['Pointer Gamut', 'Rec. 709', 'Rec. 2020', 'DCI-P3','ACES RGB'],
        standalone=False)

    alpha_p, colour_p = 0.85, 'black'

    pylab.scatter(xy[:, 0], xy[:, 1], alpha=alpha_p / 2, color=colour_p,
                  marker='+')

    display(standalone=True)


def within_boundaries(point, triangulation):
    simplex = triangulation.find_simplex(point, tol=1e-7)
    return True if simplex != -1 else False


def mask_legal_colours(image, points, xy):
    triangulation = Delaunay(points)
    image_illegal = np.copy(image).reshape((-1, 3))
    legal_colours_mask = np.array([within_boundaries(x, triangulation)
                          for x in xy])
    image_illegal[legal_colours_mask] = np.array([0, 0, 0])
    return image_illegal.reshape(image.shape)


for path in exr_images:
    message_box(os.path.basename(path))
    image = exr_image_as_array(path)

    image_plot(image)
    for key, value in colourspaces_mapping.items():
        if key in path:
            colourspace = value
            break

    message_box('Image Chromaticities')
    xy = np.array([RGB_to_xy(RGB, colourspace) for RGB in image.reshape((-1, 3))])
    chromaticities_plot(xy, colourspace)

    message_box('Rec. 709 - Illegal Colours Image')
    Rec_709_image_illegal = mask_legal_colours(image, colour.REC_709_COLOURSPACE.primaries, xy)
    image_plot(Rec_709_image_illegal)

    message_box('DCI-P3 - Illegal Colours Image')
    DCI_P3_image_illegal = mask_legal_colours(image, colour.DCI_P3_COLOURSPACE.primaries, xy)
    image_plot(DCI_P3_image_illegal)

    message_box('Rec. 2020 - Illegal Colours Image')
    Rec_2020_image_illegal = mask_legal_colours(image, colour.REC_2020_COLOURSPACE.primaries, xy)
    image_plot(Rec_2020_image_illegal)

    message_box('Pointer\'s Gamut - Illegal Colours Image')
    pointer_gamut_illegal_image = mask_legal_colours(image, colour.POINTER_GAMUT_BOUNDARIES, xy)
    image_plot(pointer_gamut_illegal_image)

    print('\n')