# import
# ------------------------------------------------------------------------------------------
import os, colour, sklearn.cluster, skimage.color, copy
import numpy as np
# local
import miam.image.Image as MIMG
import miam.processing.ColorSpaceTransform as MCST
import miam.image.ColorSpace as MICS
import miam.math.Distance as MDST
# ------------------------------------------------------------------------------------------
# MIAM project 2020
# ------------------------------------------------------------------------------------------
# author: remi.cozot@univ-littoral.fr
# ------------------------------------------------------------------------------------------
class Palette(object):
    """ 
    class Palette:
        attribute(s):
            name:       object name
            colorSpace: colorspace (colour.models.RGB_COLOURSPACES, Lab, etc.)
            nbColors:   number of colors in the Palette
            colorData:  array of pixels (np.ndarray)
            colors:     np array colors[0:nbColors,0:2]
                        sorted according to distance to black (in the palette colorSpace)
    """    
    # constructor
    def __init__(self, name, colors, colorSpace):
        self.name       = name
        self.colorSpace = colorSpace
        self.nbColors   = colors.shape[0]
        self.colors     = np.asarray(sorted(colors.tolist(), key = lambda u  : np.sqrt(np.dot(u,u))))
        # DEBUG
        # print("Palette.__init__():", self.name)

    # methods

    def build(image, nbColors, fast=True, method='kmean-Lab', **kwargs):
        # according to method
        if method == 'kmean-Lab':
            # with 'remove black or not'
            if 'removeBlack' in kwargs: removeBlack = kwargs['removeBlack']
            else:
                print('removeBlack set to True')
                removeBlack = True
            # to Lab then to Vector
            if fast: image = image.smartResize()
            imageLab = MCST.ColorSpaceTransform().compute(image,dest='Lab')
            imgLabDataVector = MIMG.Image.array2vector(imageLab.colorData)

            if removeBlack:
                # k-means: nb cluster = nbColors + 1
                kmeans_cluster_Lab = sklearn.cluster.KMeans(n_clusters=nbColors+1)
                kmeans_cluster_Lab.fit(imgLabDataVector)

                cluster_centers_Lab = kmeans_cluster_Lab.cluster_centers_
                
                # remove darkness one
                idxLmin = np.argmin(cluster_centers_Lab[:,0])                           # idx of darkness
                cluster_centers_Lab = np.delete(cluster_centers_Lab, idxLmin, axis=0)   # remove min from cluster_centers_Lab

            else:
                # k-means: nb cluster = nbColors
                kmeans_cluster_Lab = sklearn.cluster.KMeans(n_clusters=nbColors)
                kmeans_cluster_Lab.fit(imgLabDataVector)
                cluster_centers_Lab = kmeans_cluster_Lab.cluster_centers_

            colors = cluster_centers_Lab
        else:
            print('unknow palette method')
            colors = None

        return Palette('Palette_'+image.name,colors, MICS.ColorSpace.buildLab())

    def createImageOfPalette(self, colorWidth=100):
        if self.colorSpace.name =='Lab':
           cRGB = MCST.Lab_to_sRGB(self.colors, apply_cctf_encoding=True)
        elif self.colorSpace.name=='sRGB':
            cRGB = self.colors
        width = colorWidth*cRGB.shape[0]
        height=colorWidth
        # return image
        img = np.ones((height,width,3))

        for i in range(cRGB.shape[0]):
            xMin= i*colorWidth
            xMax= xMin+colorWidth
            yMin=0
            yMax= colorWidth
            img[yMin:yMax, xMin:xMax,0]=cRGB[i,0]
            img[yMin:yMax, xMin:xMax,1]=cRGB[i,1]
            img[yMin:yMax, xMin:xMax,2]=cRGB[i,2]
        # colorData, name, type, linear, colorspace, scalingFactor
        # DEBUG
        # print("createImageOfPalette:", self.name)
        return MIMG.Image(img, self.name, MIMG.imageType.imageType.SDR, False, MICS.ColorSpace.buildsRGB(),1.0)

    # create image of multiple palettes
    def createImageOfPalettes(palettes, colorWidth=100):
        # return image
        width = colorWidth*palettes[0].colors.shape[0]
        height=colorWidth*len(palettes)
        img = np.ones((height,width,3))

        for j,palette in enumerate(palettes):
            if palette.colorSpace.name =='Lab':
                cRGB = MCST.Lab_to_sRGB(palette.colors, apply_cctf_encoding=True)
            elif palette.colorSpace.name=='sRGB':
                cRGB = palette.colors

            for i in range(cRGB.shape[0]):
                xMin= i*colorWidth
                xMax= xMin+colorWidth
                yMin= j*colorWidth
                yMax= yMin+colorWidth
                img[yMin:yMax, xMin:xMax,0]=cRGB[i,0]
                img[yMin:yMax, xMin:xMax,1]=cRGB[i,1]
                img[yMin:yMax, xMin:xMax,2]=cRGB[i,2]

        return MIMG.Image(img, "palettes", MIMG.imageType.imageType.SDR, False, MICS.ColorSpace.buildsRGB(),1.0)

    # magic operators

    def __add__(self, other):
        # print("DEBUG[Palette.__add__(",self.name,",",other.name,") ]")

        # create a copy
        res = copy.deepcopy(self)

        if isinstance(other,type(self)): # both are Palette
            # check if we can add
            # 1 - color space should be the same
            if self.colorSpace.name == other.colorSpace.name: # in the sma ecolor space
                # 2 - number of colors
                if self.nbColors == other.nbColors:
                    # add
                    res.colors= self.colors + other.colors
                    res.colors = np.asarray(sorted(res.colors.tolist(), key = lambda u  : np.sqrt(np.dot(u,u))))
                else:
                    # error message
                    print("WARNING[Palette.__add__(self,other): both image palette  must have the same number of colors ! return a copy of ",self,"]")
            else:
                # error message
                print("WARNING[Palette.__add__(self,other): both image palette  must have the same color space ! return a copy of ",self,"]")
        # return 
        return res
    
    def __radd__(self, other): 
        # print("DEBUG[Palette.__radd__(",self.name,",",other.name,") ]")

        return self.__add__(other)

    def __mul__ (self, other):
        # print("DEBUG[Palette.__mul__(",self.name,",",other,") ]")

        # create a copy
        res = copy.deepcopy(self)

        if  isinstance(other, (int, float)):
            res.colors = self.colors * other
        else:
            # error message
            print("WARNING[Palette.__mul__(self,other): other must be int or float ! return a copy of ",self,"]")

        return res

    def __rmul__ (self, other): 
        # print("DEBUG[Palette.__rmul__(",self.name,",",other,") ]")

        return self.__mul__(other)