"""
@author: cuny
@file: MakeWhiter.py
@time: 2022/7/2 14:28
@description: 
美白算法
"""
import os
import cv2
import math
import numpy as np
local_path = os.path.dirname(__file__)


class MakeWhiter(object):
    class __LutWhite:
        """
        美白的内部类
        """

        def __init__(self, lut):
            cube64rows = 8
            cube64size = 64
            cube256size = 256
            cubeScale = int(cube256size / cube64size)  # 4

            reshapeLut = np.zeros((cube256size, cube256size, cube256size, 3))
            for i in range(cube64size):
                tmp = math.floor(i / cube64rows)
                cx = int((i - tmp * cube64rows) * cube64size)
                cy = int(tmp * cube64size)
                cube64 = lut[cy:cy + cube64size, cx:cx + cube64size]  # cube64 in lut(512*512 (512=8*64))
                _rows, _cols, _ = cube64.shape
                if _rows == 0 or _cols == 0:
                    continue
                cube256 = cv2.resize(cube64, (cube256size, cube256size))
                i = i * cubeScale
                for k in range(cubeScale):
                    reshapeLut[i + k] = cube256
            self.lut = reshapeLut

        def imageInLut(self, src):
            arr = src.copy()
            bs = arr[:, :, 0]
            gs = arr[:, :, 1]
            rs = arr[:, :, 2]
            arr[:, :] = self.lut[bs, gs, rs]
            return arr

    def __init__(self, lutImage: np.ndarray = None):
        self.__lutWhiten = None
        if lutImage is not None:
            self.__lutWhiten = self.__LutWhite(lutImage)

    def setLut(self, lutImage: np.ndarray):
        self.__lutWhiten = self.__LutWhite(lutImage)

    @staticmethod
    def generate_identify_color_matrix(size: int = 512, channel: int = 3) -> np.ndarray:
        """
        用于生成一张初始的查找表
        Args:
            size: 查找表尺寸，默认为512
            channel: 查找表通道数，默认为3

        Returns:
            返回生成的查找表图像
        """
        img = np.zeros((size, size, channel), dtype=np.uint8)
        for by in range(size // 64):
            for bx in range(size // 64):
                for g in range(64):
                    for r in range(64):
                        x = r + bx * 64
                        y = g + by * 64
                        img[y][x][0] = int(r * 255.0 / 63.0 + 0.5)
                        img[y][x][1] = int(g * 255.0 / 63.0 + 0.5)
                        img[y][x][2] = int((bx + by * 8.0) * 255.0 / 63.0 + 0.5)
        return cv2.cvtColor(img, cv2.COLOR_RGB2BGR).clip(0, 255).astype('uint8')

    def run(self, src: np.ndarray, strength: int) -> np.ndarray:
        """
        美白图像
        Args:
            src: 原图
            strength: 美白强度，0 - 10
        Returns:
            美白后的图像
        """
        dst = src.copy()
        strength = min(10, int(strength)) / 10.
        if strength <= 0:
            return dst
        self.setLut(cv2.imread(f"{local_path}/lut_image/3.png", -1))
        _, _, c = src.shape
        img = self.__lutWhiten.imageInLut(src[:, :, :3])
        dst[:, :, :3] = cv2.addWeighted(src[:, :, :3], 1 - strength, img, strength, 0)
        return dst


if __name__ == "__main__":
    # makeLut = MakeWhiter()
    # cv2.imwrite("lutOrigin.png", makeLut.generate_identify_color_matrix())
    input_image = cv2.imread("test_image/7.jpg", -1)
    lut_image = cv2.imread("lut_image/3.png")
    makeWhiter = MakeWhiter(lut_image)
    output_image = makeWhiter.run(input_image, 10)
    cv2.imwrite("makeWhiterCompare.png", np.hstack((input_image, output_image)))