# import
# ------------------------------------------------------------------------------------------
from . import Processing
from . import ColorSpaceTransform
from .. import image
import colour, copy 
import skimage, skimage.morphology, skimage.color, skimage.util, skimage.filters
import numpy as np
# ------------------------------------------------------------------------------------------
# MIAM project 2020
# ------------------------------------------------------------------------------------------
# author: remi.cozot@univ-littoral.fr
# ------------------------------------------------------------------------------------------
class ContrastControl(Processing.Processing):
    """description of class"""

    def __init__(self):
        pass

    def compute(self, img, **kwargs): 
        res = copy.deepcopy(img)
        if kwargs:
            if 'method' in kwargs:
                if kwargs['method'] == 'stretch':
                    if 'range' in kwargs:
                        range = kwargs['range']
                    else:
                        range = (2,98)
                     
                    p2, p98 = np.percentile(res.colorData, (2, 98))
                    res.colorData = skimage.exposure.rescale_intensity(res.colorData, in_range=(p2, p98))
                    
                elif kwargs['method'] == 'globalEqualization':
                    res.colorData = skimage.exposure.equalize_hist(res.colorData)

                elif kwargs['method'] == 'adaptativeEqualization':
                    if 'limit' in kwargs:
                        limit = kwargs['limit']
                    else:
                        limit =0.03
                    res.colorData = skimage.exposure.equalize_adapthist(res.colorData, clip_limit=limit)
                elif kwargs['method'] == 'localEqualization':
                    if 'size' in kwargs:
                        size = kwargs['size']
                    else:
                        size =min(img.shape[0],img.shape[1])

                    #print("ContrastControl.compute():'method'=localEqualization, size=",size)
                    yuv = skimage.color.rgb2yuv(res.colorData)
                    y256 = skimage.util.img_as_ubyte(yuv[:,:,0])

                    # equalization with skimage
                    selem = skimage.morphology.disk(size)
                    y256_eq = skimage.filters.rank.equalize(y256, selem=selem)
                    y_eq = skimage.util.img_as_float(y256_eq)

                    # transfert new y
                    yuv_new = copy.deepcopy(yuv)

                    y_new = yuv_new[:,:,0]
                    u_new = yuv_new[:,:,1]
                    v_new = yuv_new[:,:,2]

                    yuv_old = copy.deepcopy(yuv)
                    y_old = yuv_old[:,:,0]
                    u_old = yuv_old[:,:,1]
                    v_old = yuv_old[:,:,2]

                    y_new  = y_eq
                    y_positive = y_old>0 
                    u_new[y_positive] = y_eq[y_positive]*u_old[y_positive]/y_old[y_positive]
                    v_new[y_positive] = y_eq[y_positive]*v_old[y_positive]/y_old[y_positive]

                    yuv_new[:,:,0] = y_new
                    yuv_new[:,:,1] = u_new
                    yuv_new[:,:,2] = v_new

                    # go back to rgb
                    newRGB = skimage.color.yuv2rgb(yuv_new)
                    newRGB[newRGB>1] =1
                    newRGB[newRGB<0] = 0

                    res.colorData = newRGB
                    pass
        return res