import cv2
import numpy as np
import matplotlib.pyplot as plt
from vispy.color import (get_colormap, get_colormaps, Colormap)


class Render:
    ''' pass '''
    def __init__(self):
        pass
#         self.load = False

    def load_brush(self, brush):
        '''
            brush: 4 channel array
        '''
        self.B, self.G, self.R, self.A = cv2.split(brush)
        self.a_trans = self.A

    def load_A_channel(self, A):
        self.A = A
        self.a_trans = self.A

    def show_img(self):

        plt.imshow(cv2.merge((self.R, self.G, self.B)))
        plt.show()

    def emboss(self, img, intensity):
        '''input: 
        	img: 4D BGRA
        	The intensity of emboss
        '''
        img = cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)
        (H, L, S) = cv2.split(cv2.cvtColor(img, cv2.COLOR_RGB2HLS))

        # Normalisation：For HSV, Hue range is [0,179], Saturation range is [0,255] and Value range is [0,255].
        H, L, S = H / 179., L / 255., S / 255.
        H, L, S = H, L, S
        # Apply the convolution kernel to channel L：
        p, q, s, t = 0., 0., intensity, 0.
        # p,q,s control the intensity of the relief in three directions
        # The value of matrix center controls the overall brightness, here sets 1.0
        kernel = np.array([[-q, -s, t], [-p, 1.0, p], [-t, s, q]],
                          dtype=np.float32)
        #         kernel = np.array([[ 0, -1,-2],
        #                            [ 1, 0, -1],
        #                            [ 2,  1, 0]], dtype=np.float32)

        L = cv2.filter2D(L, -1, kernel)
        L = np.clip(L, 0, 1)
        H, L, S = H * 179., L * 255., S * 255.

        H, L, S = H.astype('uint8'), L.astype('uint8'), S.astype('uint8')
        #-------------------------------------------------------------------
        # If in the previous step makes the brightness too large, please re-limit the brightness here
        # H,L,S = H, Gamma_trans(L, 255, 1.2), Gamma_trans(S, 255, 1)
        # H,L,S = H, MaxMinNormal(L, 1, 150), MaxMinNormal(S, 0,255)
        #-------------------------------------------------------------------
        # Merge each channel
        res = cv2.merge((H, L, S))
        res = cv2.cvtColor(res, cv2.COLOR_HLS2RGB)
        return cv2.cvtColor(res, cv2.COLOR_RGB2BGRA)

    @classmethod
    def colorMapInterp(cls, colors, N):
        '''
        colors: a list of RGB color array/list/tuple
                For example: colors = ([1,0,0],[0,1,0],...) 
                or colors = ('r','g') 
        N: number of interpolation 
        return: rgb interpolation array. type: float [0,1]
        '''
        rm = Colormap(colors)
        rgb_interp = rm[np.linspace(0, 1, N)].rgb  #* 255
        # show colors
        # plt.imshow(rgb_interp.reshape(1,N,3).astype('uint8'))
        # plt.axis('off')
        return rgb_interp  # float [0,1]

    @classmethod
    def ramp(cls, c1, c2, c3, c4, h0=None, w0=None):
        '''
        input:
            c: Color of four vertices
        return:
            2D color bgr map array. type: float [0,1]
        '''
        if h0 == None:
            h0, w0 = cls.A.shape
        # scale = 20
        # h, w =  h0//scale+1, w0//scale+1 # h, w can be replace others and then we interpolate the output image
        h, w = 20, 20
        rec1 = cls.colorMapInterp((c1, c2), h)

        rec3 = cls.colorMapInterp((c3, c4), h)

        color_mat = np.zeros((h, w, 3))
        for i in range(h):
            color_mat[i, :, :] = cls.colorMapInterp(
                (rec1[i, :], rec3[h - i - 1, :]), w)
        # show it
        # plt.imshow(color_mat)
        color_mat[:, :, [0, 2]] = color_mat[:, :, [2, 0]]  # rgb -> bgr

        return cv2.resize(color_mat, (w0, h0))  # float [0,1]

    @staticmethod
    def get_random_color(bgr, sigma):
        '''
        rgb: [b, g, r] float value
        sigma: [sigma_b, sigma_g, sigma_r] float range
        return rgb color
        '''
        #         color = [np.random.normal(bgr[2], sigma[0]), np.random.normal(bgr[1], sigma[1]), np.random.normal(bgr[0], sigma[2])]
        color = [
            bgr[2] + sigma[0] * (np.random.rand() - 1 / 2),
            bgr[1] + sigma[1] * np.random.rand(),
            bgr[0] * sigma[2] + np.random.rand()
        ]
        return np.clip(color, 0, 1)

    @classmethod
    def gen_colormap(cls, h, w, c1, sigma=[0.5, 0.5, 0.1]):
        '''
        c1: [b, g, r] float value
        sigma: [sigma_r, sigma_g, sigma_b] float range
        '''
        c2 = Render.get_random_color(c1, sigma)
        c3 = Render.get_random_color(c1, sigma)
        c4 = Render.get_random_color(c1, sigma)
        color = Render.ramp(c1, c2, c3, c4, h, w)
        return color

    def get_fore_back(self, canvas, brush):
        '''
        Get foreground and background of canvas and brush
        '''
        # brush = cv2.bilateralFilter(brush,15, 35,35)
        H, W = canvas.shape[:2]  # canvas shape
        h_brush, w_brush = brush.shape[:2]

        h_coef = h_brush / w_brush  # keep the same shape of brush
        scale = 0.4  # scale coefficient for roi
        h, w = int(H * h_coef * scale), int(W * scale)  # roi shape
        # Attention! cv2.resize(img, (w,h)) NOT (h, w)
        brush = cv2.resize(brush, (w, h), interpolation=cv2.INTER_AREA)
        #         plt.figure(figsize=(12,12))
        #         plt.subplot(231)
        #         plt.imshow(cv2.cvtColor(brush, cv2.COLOR_BGR2RGB)); plt.title('brush(padding)')
        # padding
        brush = cv2.copyMakeBorder(brush,
                                   int(np.ceil((H - h) / 2)),
                                   int(np.floor((H - h) / 2)),
                                   int(np.ceil((W - w) / 2)),
                                   int(np.floor((W - w) / 2)),
                                   cv2.BORDER_CONSTANT,
                                   value=[255, 255, 255])
        #         print(canvas.shape, brush.shape)
        #         plt.figure(figsize=(12,12))
        #         plt.subplot(231)
        #         plt.imshow(cv2.cvtColor(brush, cv2.COLOR_BGR2RGB)); plt.title('brush(padding)')

        brush = cv2.bilateralFilter(
            brush, 15, 35,
            35)  # this step can be delete if input image is 'clean'.
        brushgray = cv2.cvtColor(brush, cv2.COLOR_BGR2GRAY)
        ret, mask = cv2.threshold(
            brushgray, 250, 255,
            cv2.THRESH_BINARY)  #ret is thresh（250）mask is Binary image

        mask_inv = cv2.bitwise_not(mask)
        #plt.subplot(232);plt.imshow(mask);plt.title('mask')
        #plt.subplot(233);plt.imshow(mask_inv);plt.title('mask_inv')

        canvas_bg = cv2.bitwise_and(
            canvas, brush,
            mask=mask)  # 将mask去掉，可以看出 _and操作是取数值更低（更黑）的元素,而mask是将mask对应元素置于0
        brush_bg = cv2.bitwise_and(canvas, canvas, mask=mask_inv)
        #plt.subplot(234);plt.imshow(cv2.cvtColor(canvas_bg, cv2.COLOR_BGR2RGB));plt.title('canvas_brush_mask_and')
        #plt.subplot(235);plt.imshow(cv2.cvtColor(brush_bg, cv2.COLOR_BGR2RGB));plt.title('brush_brush_mask_inv_and')

        brush_fg = cv2.bitwise_and(brush, brush, mask=mask_inv)
        #plt.subplot(236);plt.imshow(cv2.cvtColor(brush_fg, cv2.COLOR_BGR2RGB));plt.title('brush_fg')

        return canvas_bg, brush_bg, brush_fg  # dtype: uint8

    def rotate_origin_img(self, img, theta):
        '''
        theta: Counterclockwise Angle
        This function is going to rotate and pad the image to the right size
        '''
        h_origin, w_origin = img.shape[:2]

        # padding
        size = int(np.ceil(np.sqrt(w_origin**2 + h_origin**2)))
        top = int(np.ceil((size - h_origin) / 2))
        bottom = int(np.ceil((size - h_origin) / 2))
        left = int(np.ceil((size - w_origin) / 2))
        right = int(np.ceil((size - w_origin) / 2))
        padding_img = cv2.copyMakeBorder(img,
                                         top,
                                         bottom,
                                         left,
                                         right,
                                         cv2.BORDER_CONSTANT,
                                         value=[0, 0, 0])

        rotate_origin = ((size - 1) / 2, (size - 1) / 2)
        M = cv2.getRotationMatrix2D(rotate_origin, theta, 1)
        rotated_img = cv2.warpAffine(padding_img,
                                     M, (size, size),
                                     flags=cv2.INTER_LANCZOS4)

        assert rotated_img.dtype == 'uint8'
        return rotated_img, dict([('h_origin', h_origin),
                                  ('w_origin', w_origin), ('theta', theta),
                                  ('top', top), ('bottom', bottom),
                                  ('left', left), ('right', right),
                                  ('rotate_origin', rotate_origin)])

    def inv_rotate_padding_img(self, padding_img, dic):

        h, w = padding_img.shape[:2]
        M_inv = cv2.getRotationMatrix2D(dic['rotate_origin'], -dic['theta'], 1)
        inv_rot_padding_img = cv2.warpAffine(padding_img,
                                             M_inv, (h, w),
                                             flags=cv2.INTER_LANCZOS4)

        # clip
        inv_rot_img = inv_rot_padding_img[dic['top'] - 1:-dic['bottom'],
                                          dic['left'] - 1:-dic['right'],
                                          ]  # if dimension>2, add ':'

        inv_rot_img = cv2.resize(inv_rot_img,
                                 (dic['w_origin'], dic['h_origin']),
                                 interpolation=cv2.INTER_LANCZOS4)

        assert inv_rot_img.shape[:2] == (dic['h_origin'], dic['w_origin'])
        assert inv_rot_img.dtype == 'uint8'
        return inv_rot_img

    def transp_mat(self, size, alpha_min, alpha_max, rate):
        '''
        Args:
            size: rotated image shape
        give the transparent matrix.
        '''
        N = 50
        mat = np.ones((N, N))
        x = np.linspace(1 - alpha_max, 1 - alpha_min, N)
        alpha = ((1 - utils.Gamma_trans(x, rate)) * (alpha_max - alpha_min) +
                 alpha_min)
        for i in range(N):
            mat[i, :] = alpha[i] * mat[i, :]
        mat = cv2.resize(mat, (size))
        return mat

    def transparency_rotate(self, theta=0, alpha_min=0.3, alpha_max=1, rate=1):
        '''
        Apply transparency to channel A.
        Args
            a: channel A
            
        '''
        a = self.A
        a, coord_dic = self.rotate_origin_img(a, theta)
        a_trans = a * self.transp_mat(a.shape, alpha_min, alpha_max, rate)
        a_trans = self.inv_rotate_padding_img(a_trans.astype('uint8'),
                                              coord_dic)
        self.a_trans = a_trans

    def transparency_love(self, love_mat):
        h, w = self.A.shape
        love_mat = cv2.resize(love_mat, (w, h)) * 255
        self.a_trans = love_mat.astype('uint8')

    def overlay_with_transparency(self, bgimg, xmin=0, ymin=0):
        '''
        bgimg: a 4 channel image, use as background
        xmin, ymin: a corrdinate in bgimg. from where the fgimg will be put
        trans_percent: transparency of fgimg. [0.0,1.0]
        '''
        trans_percent = 1
        fgimg = cv2.merge((self.B, self.G, self.R, self.a_trans))

        #create roi and black mask
        roi = bgimg[ymin:ymin + fgimg.shape[0], xmin:xmin + fgimg.shape[1]]
        roi_black = np.zeros(fgimg.shape).astype(bgimg.dtype)

        #convert alpha to [0,1] and reshape it so that we can to math operation later
        alpha = np.expand_dims((self.a_trans.astype('float32') / 255.),
                               axis=-1)
        alpha_inv = 1. - alpha

        roi_bg_black = (roi * alpha_inv).astype(bgimg.dtype)
        roi_bg_trans = cv2.addWeighted(roi, 1 - trans_percent, roi_bg_black,
                                       trans_percent, 0)

        fg_black = (fgimg * alpha).astype(fgimg.dtype)
        fg_trans = cv2.addWeighted(roi_black, 1 - trans_percent, fg_black,
                                   trans_percent, 0)

        result = bgimg.copy()
        result[ymin:ymin + fgimg.shape[0],
               xmin:xmin + fgimg.shape[1]] = cv2.add(fg_trans, roi_bg_trans)

        return result
# -------------------------------------------------------------

    def over_layer(self, img1, img2, transp1=1, transp2=1):
        '''
        input: float bgra
            img1: foreground   img2: background
            transp1: foreground transparency
        output: float bgra
        '''
        b1, g1, r1, alpha1 = cv2.split(img1)
        b2, g2, r2, alpha2 = cv2.split(img2)
        # ----------
        alpha1 = transp1 * alpha1
        alpha2 = transp2 * alpha2
        # ----------
        alpha = alpha1 + alpha2 * (1 - alpha1)
        alpha = np.clip(alpha, 0, 1)

        def outRGB(c1, c2):
            c1[alpha == 0] = np.nan
            c2[alpha == 0] = np.nan
            c = (c1 * alpha1 + c2 * alpha2 * (1 - alpha1)) / alpha
            return np.nan_to_num(c)

        b = outRGB(b1, b2)
        g = outRGB(g1, g2)
        r = outRGB(r1, r2)
        b, g, r = np.clip(b, 0, 1), np.clip(g, 0, 1), np.clip(r, 0, 1)
        # b,g,r,alpha = utils.append1D(b,g,r,alpha)
        # bgra = np.concatenate([b,g,r,alpha], axis=2)
        bgra = cv2.merge([b, g, r, alpha])
        return bgra

    def overlay(self,
                img1,
                img2,
                xmin=0,
                ymin=0,
                transp_brush=1,
                transp_canvas=1):
        '''input: 
                float img1: foreground  img2: background
           return: float overlay
        '''
        roi = img2[ymin:ymin + img1.shape[0], xmin:xmin + img1.shape[1]]
        roi = self.over_layer(img1,
                              roi,
                              transp1=transp_brush,
                              transp2=transp_canvas)
        img2[ymin:ymin + img1.shape[0], xmin:xmin + img1.shape[1]] = roi
        return img2


# -------------------------------------------------------------

    @staticmethod
    def zeroAlpha2zeroRGB(img):  # this is not used
        '''
        input: 
            img: bgra 4d array, float
        return: bgra 4d array, float
        '''
        b, g, r, a = cv2.split(img)
        r[a == 0] = 0
        g[a == 0] = 0
        b[a == 0] = 0

        r, g, b, a = utils.append1D(r, g, b, a)
        return np.concatenate([b, g, r, a], axis=2)


class utils:
    @staticmethod
    def append1D(*array_list):
        ''' For exemple: 
                before: array.shape (5,2)
                after: array.shape  (5,2,1)
            '''
        out = []
        for array in array_list:
            out.append(array.reshape(array.shape + (1, )))
        if len(out) == 1:
            return out[0]
        else:
            return [i for i in out]

    @staticmethod
    def myplot(img, channel='bgr', figsize=6):
        ''' can show image channel(float/uint8) BGR,BGRA,HLS,HSV.
        if RGB, set channel = 0'''
        def cvt2rgb(img, channel='bgr'):
            '''it can convert image channel BGR,BGRA,HLS,HSV to RGB'''
            if channel == 'bgr' or channel == 'BGR':
                print('bgr', img.shape)
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            elif channel == 'bgra' or channel == 'BGRA':
                print('bgra', img.shape)
                #                 b,g,r,a = cv2.split(img)
                #                 r,g,b,a = utils.append1D(r,g,b,a)
                #                 img = np.concatenate([r,g,b,a], axis=2)
                #                 img = np.c_[rgb, np.reshape(img[:,:,-1],rgb.shape[:2] +(1,))]
                img[:, :, [2, 0]] = img[:, :, [0, 2]]
            elif channel == 'HLS' or channel == 'hls':
                img = cv2.cvtColor(img, cv2.COLOR_HLS2RGB)
                print('hls', img.shape)
            elif channel == 'HSV' or channel == 'hsv':
                img = cv2.cvtColor(img, cv2.COLOR_HSV2RGB)
                print('hsv', img.shape)
            else:
                return img
            return img

        if img.dtype == np.float or img.dtype == np.float32:
            img = np.uint8(img * 255)
        if channel:
            img = cvt2rgb(img, channel)
        plt.figure(figsize=(figsize, figsize))
        plt.imshow(img)
        plt.show()

    @staticmethod
    def Gamma_trans(I, gamma):
        '''
        param:
            gamma: if your intersted region is too bright, set gamma > 1 decreasing contrast.
               and if your intersted region is too dark, set 1> gamma > 0 to increase contrast.
            I_max: is the maximun of the channel of I.
        return:
            the map of I
        '''
        I_max = I.max()
        fI = I / I_max
        out = np.power(fI, gamma)
        out = out * I_max
        return out


if __name__ == '__main__':

    def test_overlay():
        canvas = cv2.imread('pic/brush_blue.png', -1)
        brush = cv2.imread('pic/pika_roi.png', -1)
        render = Render()
        canvas_ = render.overlay(brush / 255,
                                 canvas / 255,
                                 xmin=0,
                                 ymin=0,
                                 transp_brush=0.7)
        utils.myplot(canvas_, 'bgra')

    test_overlay()

    def test_ramp():
        render = Render()
        res = render.ramp('r', 'g', 'b', 'g')

    # =================================================================
#     render.load_brush(brush)
#     render.transparency_rotate(theta=0, alpha_min=0.2, alpha_max=0.99,rate=2)

#     res = render.overlay_with_transparency(canvas)
#     plt.imshow(cv2.cvtColor(res,cv2.COLOR_BGRA2RGB));plt.show()
#     res = render.emboss(res, 1)
#     plt.imshow(cv2.cvtColor(res,cv2.COLOR_BGRA2RGB));plt.show()
# =================================================================
