#!/usr/bin/env python
from __future__ import division, print_function
import os
import numpy as np
import cv2
from PIL import Image as PILImage
import glob
import matplotlib.pyplot as plt
import colour
import colour_checker_detection
from colour_checker_detection import (
    EXAMPLES_RESOURCES_DIRECTORY,
    colour_checkers_coordinates_segmentation,
    detect_colour_checkers_segmentation,
    extract_colour_checkers_segmentation)

from colour_checker_detection.detection.segmentation import (
    adjust_image, swatch_masks)


colour.utilities.describe_environment();


image_dir = '/mnt/data2tb/libraries/adapt/adapt_ros_ws/data/2022-04-21-14-24-57/images'
save_fname = '/mnt/data2tb/libraries/adapt/adapt_ros_ws/data/color_cal_mat.txt'


fnames = glob.glob('%s/*.jpg' % image_dir)

n = 0
for fname in fnames:
    if n == 0:
        image = cv2.imread(fname)[:, :, ::-1].astype(float)
    else:
        image += cv2.imread(fname)[:, :, ::-1].astype(float)

    n += 1

image /= n*255
#image = np.round(image).astype(np.uint8)
#image = cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)


if False:
    fname = '/home/local/KHQ/matt.brown/Downloads/test.png'
    fname = '/mnt/data2tb/libraries/adapt/adapt_ros_ws/data/2022-04-21-14-24-57/colorchecker/1650579897669061.jpg'
    fname = '/mnt/data2tb/libraries/adapt/adapt_ros_ws/data/2022-04-27-10-24-29/color_checker/1651083869098989.jpg'
    #image = cv2.imread(fname)[:, :, ::-1]
    image = colour.io.read_image(fname)[:, :, :3]


image0 = image
image = colour.cctf_decoding(image0)


ret = detect_colour_checkers_segmentation(image, additional_data=True)[0]
swatches, colour_checker_image, masks = ret


# Using the additional data to plot the colour checker and masks.
masks_i = np.zeros(colour_checker_image.shape)
for i, mask in enumerate(masks):
    masks_i[mask[0]:mask[1], mask[2]:mask[3], ...] = 1

colour.plotting.plot_image(
    colour.cctf_encoding(
        np.clip(colour_checker_image + masks_i * 0.25, 0, 1)));


D65 = colour.CCS_ILLUMINANTS['CIE 1931 2 Degree Standard Observer']['D65']

if False:
    REFERENCE_COLOUR_CHECKER = colour.CCS_COLOURCHECKERS[
        'ColorChecker24 - After November 2014']
else:
    # Spyder24
    from colour.characterisation.datasets.colour_checkers.chromaticity_coordinates import ColourChecker
    from collections import OrderedDict

    #REFERENCE_COLOUR_CHECKER = colour.CCS_COLOURCHECKERS['cc2005']
    #REFERENCE_COLOUR_CHECKER = colour.CCS_COLOURCHECKERS['ccb2014']
    REFERENCE_COLOUR_CHECKER = colour.CCS_COLOURCHECKERS['cca2014']
    order = [5, 4, 3, 2, 1, 0, 6, 7, 8, 9, 10, 11, 17, 16, 15, 14, 13, 12, 18,
             19, 20, 21, 22, 23]
    keys = list(REFERENCE_COLOUR_CHECKER.data.keys())
    keys = [keys[i] for i in order]
    data = OrderedDict([(key, REFERENCE_COLOUR_CHECKER.data[key]) for key in keys])
    REFERENCE_COLOUR_CHECKER = ColourChecker('Spyder24', data, illuminant=D65)



REFERENCE_SWATCHES = colour.XYZ_to_RGB(
        colour.xyY_to_XYZ(list(REFERENCE_COLOUR_CHECKER.data.values())),
        REFERENCE_COLOUR_CHECKER.illuminant, D65,
        colour.RGB_COLOURSPACES['sRGB'].matrix_XYZ_to_RGB)


swatches_xyY = colour.XYZ_to_xyY(colour.RGB_to_XYZ(
        swatches, D65, D65, colour.RGB_COLOURSPACES['sRGB'].matrix_RGB_to_XYZ))

colour_checker = colour.characterisation.ColourChecker('test',
        dict(zip(REFERENCE_COLOUR_CHECKER.data.keys(), swatches_xyY)),
        D65)

colour.plotting.plot_multi_colour_checkers(
    [REFERENCE_COLOUR_CHECKER, colour_checker])

swatches_f = colour.colour_correction(swatches, swatches, REFERENCE_SWATCHES,
                                      degree=4)
swatches_f_xyY = colour.XYZ_to_xyY(colour.RGB_to_XYZ(
    swatches_f, D65, D65, colour.RGB_COLOURSPACES['sRGB'].matrix_RGB_to_XYZ))
colour_checker = colour.characterisation.ColourChecker('test',
    dict(zip(REFERENCE_COLOUR_CHECKER.data.keys(), swatches_f_xyY)),
    D65)

colour.plotting.plot_multi_colour_checkers(
    [REFERENCE_COLOUR_CHECKER, colour_checker])


if False:
    # Apply an intensity correction to best match the intensity of the colors.
    plt.plot(swatches[:, 2], REFERENCE_SWATCHES[:, 2], 'ro')


corrected = colour.colour_correction(image, swatches, REFERENCE_SWATCHES,
                                     degree=4)
#corrected -= corrected.min()
#corrected /= corrected.max()
corrected = np.clip(corrected, 0, 1)
srgb = colour.cctf_encoding(corrected)
srgb = np.clip(srgb, 0, 1)
#plt.imshow(srgb);
colour.plotting.plot_image(srgb)
#colour.plotting.plot_image(colour.cctf_encoding(image))


M = colour.matrix_colour_correction(swatches, REFERENCE_SWATCHES, degree=4)
image_ = np.reshape(image, (image.shape[0]*image.shape[1], 3)).T
image_ = np.dot(M, image_)
image_ = np.reshape(image_.T, (image.shape[0], image.shape[1], 3))
colour.plotting.plot_image(colour.cctf_encoding(image_))


def srgb_gamma_encode(image):
    ind = image > 0.0031308
    image[ind] = 1.055*image[ind]**(1/2.4) - 0.055
    image[~ind] = 12.92*image[~ind]
    return image


def srgb_gamma_decode(image):
    dtype = image.dtype
    if dtype == np.uint8:
        image = image.astype(float)/255
    elif dtype == np.uint16:
        image = image.astype(float)/65535

    ind = image > 0.04045
    image[ind] = ((image[ind] + 0.055)/1.055)**2.4
    image[~ind] = image[~ind]/12.92

    return image



print('M = [[%0.10f, %0.10f, %0.10f],\n     [%0.10f, %0.10f, %0.10f],\n     [%0.10f, %0.10f, %0.10f]]' %
      tuple(M.ravel()))

np.savetxt(save_fname, M)
