# 功能函数
import os
import cv2
import numpy as np
from PIL import Image, ImageFilter
from matplotlib import pyplot as plt
import collections
import colorsys

#定义字典存放颜色分量上下限
#例如：{颜色: [min分量, max分量]}
#{'red': [array([160, 43, 46]), array([179, 255, 255])]}  
def getColorList():
    dict = collections.defaultdict(list)

    # 黑色
    lower_black = np.array([0, 0, 0])
    upper_black = np.array([180, 255, 46])
    color_list = []
    color_list.append(lower_black)
    color_list.append(upper_black)
    dict['black'] = color_list

    #灰色
    lower_gray = np.array([0, 0, 46])
    upper_gray = np.array([180, 43, 220])
    color_list = []
    color_list.append(lower_gray)
    color_list.append(upper_gray)
    dict['gray']=color_list

    # 白色
    lower_white = np.array([0, 0, 221])
    upper_white = np.array([180, 30, 255])
    color_list = []
    color_list.append(lower_white)
    color_list.append(upper_white)
    dict['white'] = color_list
    
    # 红色1
    lower_red = np.array([0, 43, 46])
    upper_red = np.array([10, 255, 255])
    color_list = []
    color_list.append(lower_red)
    color_list.append(upper_red)
    dict['red'] = color_list
    
    #红色2
    lower_red = np.array([156, 43, 46])
    upper_red = np.array([180, 255, 255])
    color_list = []
    color_list.append(lower_red)
    color_list.append(upper_red)
    dict['red2']=color_list

    #橙色
    lower_orange = np.array([11, 43, 46])
    upper_orange = np.array([25, 255, 255])
    color_list = []
    color_list.append(lower_orange)
    color_list.append(upper_orange)
    dict['orange'] = color_list

    #黄色
    lower_yellow = np.array([26, 43, 46])
    upper_yellow = np.array([34, 255, 255])
    color_list = []
    color_list.append(lower_yellow)
    color_list.append(upper_yellow)
    dict['yellow'] = color_list

    #绿色
    lower_green = np.array([35, 43, 46])
    upper_green = np.array([77, 255, 255])
    color_list = []
    color_list.append(lower_green)
    color_list.append(upper_green)
    dict['green'] = color_list

    #青色
    lower_cyan = np.array([78, 43, 46])
    upper_cyan = np.array([99, 255, 255])
    color_list = []
    color_list.append(lower_cyan)
    color_list.append(upper_cyan)
    dict['cyan'] = color_list

    #蓝色
    lower_blue = np.array([100, 43, 46])
    upper_blue = np.array([124, 255, 255])
    color_list = []
    color_list.append(lower_blue)
    color_list.append(upper_blue)
    dict['blue'] = color_list

    # 紫色
    lower_purple = np.array([125, 43, 46])
    upper_purple = np.array([155, 255, 255])
    color_list = []
    color_list.append(lower_purple)
    color_list.append(upper_purple)
    dict['purple'] = color_list

    return dict


# 输入hsv图片得到颜色

def get_color_from_pos(x, y, hsv, color_dict):
    
    assert x >= 0 and x < hsv.shape[1]
    assert y >= 0 and y < hsv.shape[0]
    
    point_hsv = hsv[int(y), int(x),:]
    h, s, v = point_hsv[0], point_hsv[1], point_hsv[2]
    for c in color_dict:
        low_rnk, high_rnk = color_dict[c][0], color_dict[c][1]
        if h >= low_rnk[0] and h <= high_rnk[0] and s >= low_rnk[1] and s <= high_rnk[1] and v >= low_rnk[2] and v <= high_rnk[2]:
            return c, point_hsv
    
    
def get_color_from_image(hsv, color_dict):
    h, w, _ = hsv.shape
    center_hsv = hsv[int(h/2-1), int(w/2-1),:]
    h, s, v = center_hsv[0], center_hsv[1], center_hsv[2]
    print('source hsv', h, s, v)
    for c in color_dict:
        low_rnk, high_rnk = color_dict[c][0], color_dict[c][1]
        if h >= low_rnk[0] and h <= high_rnk[0] and s >= low_rnk[1] and s <= high_rnk[1] and v >= low_rnk[2] and v <= high_rnk[2]:
            return c, center_hsv

def get_color_from_image_2(hsv, color_dict):
    hist, bins = np.histogram(hsv[:, :, 0], bins=180, range=(0,180))
    hist[0] = 0
    for i in range(len(hist)):
        if hist[i] == hist.max():
            idx = i
            break
    print(idx)
    for c in color_dict:
        low_rnk, high_rnk = color_dict[c][0], color_dict[c][1]
        if idx >= low_rnk[0] and idx <= high_rnk[0]:
            return c, idx
    
        
# 输入rgb值得到hsv值
def rgb2hsv(rgb, color_dict):
    assert len(rgb) == 3
    r, g, b = rgb[0], rgb[1], rgb[2]
    hsv = colorsys.rgb_to_hsv(r/255, g/255, b/255)
    hsv = np.array(hsv) * np.array([180, 255, 255])
    h, s, v = int(hsv[0]), int(hsv[1]), int(hsv[2])
    print('target hsv', h, s, v)
    
    for c in color_dict:
        low_rnk, high_rnk = color_dict[c][0], color_dict[c][1]
        if h >= low_rnk[0] and h <= high_rnk[0] and s >= low_rnk[1] and s <= high_rnk[1] and v >= low_rnk[2] and v <= high_rnk[2]:
            return c, hsv

# 判断颜色是否为彩色，黑白灰返回false
def judge_color(color):
    if 'black' in color or 'gray' in color or 'white' in color:
        return False
    else:
        return True

    
# 获得遮罩
def get_mask(color, hsv, color_dict):
    if 'red' in color:
        mask1 = cv2.inRange(hsv, color_dict['red'][0], color_dict['red'][1])
        mask2 = cv2.inRange(hsv, color_dict['red2'][0], color_dict['red2'][1])
        mask = mask1 + mask2
    else:
        mask = cv2.inRange(hsv, color_dict[color][0], color_dict[color][1])
        
    return mask != 0

# 保存mask
def save_mask(path, mask, cnt):
    mask[mask!=0] = 255
    mask = Image.fromarray(mask).convert('1')
    mask_name = os.path.join(path, str(cnt)+'.png')
    mask.save(mask_name)


# 获得局部区域的属于color范围的遮罩
def get_region_mask(color, hsv, x1, y1, x, y, color_dict):
    if 'red' in color:
        mask1 = cv2.inRange(hsv, color_dict['red'][0], color_dict['red'][1])
        mask2 = cv2.inRange(hsv, color_dict['red2'][0], color_dict['red2'][1])
        mask = mask1 + mask2
    else:
        mask = cv2.inRange(hsv, color_dict[color][0], color_dict[color][1])
    
    mask = mask != 0
    tmp = np.zeros(mask.shape)
    tmp[y1:y, x1:x] = mask[y1:y, x1:x]
    
    return tmp != 0

def get_region_mask_rect(color, hsv, x1, y1, x, y, color_dict):
    
    mask = cv2.inRange(hsv, color_dict[color][0], color_dict[color][1])
    
    mask = mask == 0 # 除了这个颜色，其他的都置为True
    tmp = np.zeros(mask.shape)
    tmp[y1:y, x1:x] = mask[y1:y, x1:x]
        
    return tmp != 0

def get_region_mask_circle(color, hsv, black_back, color_dict):
    mask = cv2.inRange(hsv, color_dict[color][0], color_dict[color][1])
    mask = mask == 0 # 除了这个颜色，其他的都置为True
    
    black_back = black_back!=0
    
    mask = mask & black_back # 不是这个颜色且在这个区域
    
    return mask

def get_region_mask_remove(color, hsv, white_back, color_dict): # 去除涂抹区域中的目标颜色
    white_back = white_back != 0 # 涂抹的区域
    return white_back

# 加光照函数 输入rgb 输出rgb
def adding_light(img, light_conf):
    rows, cols = img.shape[:2]
    
    cX_list = [0, 0, 0, rows-1, rows-1, rows-1]
    cY_list = [0, cols // 2, cols - 1, 0, cols // 2, cols - 1]
    
    radius = min(rows, cols)
    strength = light_conf['light_strength']
    cX = cX_list[light_conf['light_type']]
    cY = cY_list[light_conf['light_type']]
    
    dist = np.zeros((rows, cols))
    x = np.arange(0, rows)
    dx = np.power((x - cX), 2)

    for i in range(cols):
        y = np.zeros(rows) + i
        dy = np.power((y - cY), 2)
        dist[:, i] = dx + dy

           
    R, G, B = img[:, :, 0], img[:, :, 1], img[:, :, 2]
    res = strength * (1.0 - np.sqrt(dist) / radius)
    
    mape = dist < radius * radius
    
    R = (R  + res) * mape + R * (1 - mape)
    G = (G  + res) * mape + G * (1 - mape)
    B = (B  + res) * mape + B * (1 - mape)
        
    R = np.clip(R, 0, 255)
    G = np.clip(G, 0, 255)
    B = np.clip(B, 0, 255)
    
    res = np.uint8(np.dstack((R, G, B)))
    
    return res

def Sort_map(result):
    idx = 0
    res = {}
    for key in result.keys():
        if "best" in key:
            res[str(idx)+'.png'] = result[key]
            idx += 1
    
    for key in result.keys():
        if "best" not in key:
            res[str(idx)+'.png'] = result[key]
            idx += 1
    return res

def get_name(flag_h, flag_s, flag_v):
    save_name = ""
    if flag_h != None:
        save_name += "h_"
        save_name += str(int(flag_h[0]))
        save_name += "_"
        save_name += str(int(flag_h[1]))
        save_name += "_"
    
    if flag_s != None:
        save_name += "s_"
        save_name += str(int(flag_s[0]))
        save_name += "_"
        save_name += str(int(flag_s[1]))
        save_name += "_"
    
    if flag_v != None:
        save_name += "v_"
        save_name += str(int(flag_v[0]))
        save_name += "_"
        save_name += str(int(flag_v[1]))

    save_name = save_name.strip('_')
    return save_name + '.jpg'

def light_map(x, l_rnk, h_rnk):
    if x <= l_rnk:
        res = 0
    elif  x >= h_rnk:
        res = 255
    else:
        k = 255 / (h_rnk - l_rnk)
        b = -255 * l_rnk / (h_rnk - l_rnk)
        
        res = k * x + b
    return int(res)

def curve(hsv, mask, h_rnk, l_rnk):
    mask = np.array(mask)
    mask = mask > 10
    
    v = hsv[:, :, 2]
    v_ = v.copy()
    
    vfunc = np.vectorize(light_map, excluded=['l_rnk', 'h_rnk'])
    v_ = vfunc(v_, l_rnk=l_rnk, h_rnk=h_rnk)
    
    hsv[:, :, 2] = v_ * mask + v * (~mask)
    return hsv

# 保存结果
def save_result(hsv, alpha, mask_ori, pic_path, light_conf, curve_conf):
    if curve_conf['if_curve'] == True:
        hsv = curve(hsv=hsv, mask=mask_ori, h_rnk=curve_conf['h_rnk'], l_rnk=curve_conf['l_rnk'])
            
    if alpha == None: # 三通道图片
        tag_img = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
        # 
        # tag_img= cv2.GaussianBlur(tag_img, (9, 9), 0)
        if light_conf['if_light'] != False:
            tag_img = adding_light(tag_img, light_conf)
            
        res = Image.open(pic_path)
        res.paste(Image.fromarray(tag_img), mask = mask_ori)
    else:
        tag_img = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
        
        if light_conf['if_light'] != False:
            tag_img = adding_light(tag_img, light_conf)
        
        if mask_ori!= None:
            res = Image.open(pic_path)
            res.paste(Image.fromarray(tag_img), mask=mask_ori)
            r, g, b, _ = res.split()
            res = Image.merge('RGBA', (r, g, b, alpha))        
        else:
            r, g, b = Image.fromarray(tag_img[:, :, 0]), Image.fromarray(tag_img[:, :, 1]), Image.fromarray(tag_img[:, :, 2])
            res = Image.merge('RGBA', (r, g, b, alpha))
            
    # res.save(save_name)
    return res
        
def mapping(mape, a, l_flag, h_flag, pre_set):
    if pre_set != None:
        assert len(pre_set) == 2
        l_flag = pre_set[0]
        h_flag = pre_set[1]

    mi = mape[a].min()
    mx = mape[a].max()
    # 相交
    if (l_flag < mx and l_flag > mi) or (h_flag > mi and h_flag < mx):
        if mx > h_flag: 
            if mi >= l_flag: # [mi, h_flag]
                mape = (mape - mi) / (mx - mi) * (h_flag - mi) + mi
            else: # [l_flag, h_flag]
                mape = (mape - mi) / (mx - mi) * (h_flag - l_flag) + l_flag
        
        else:
            if mi < l_flag: # [l_flag, mx]
                mape = (mape - mi) / (mx - mi) * (mx - l_flag) + l_flag
    # 不相交
    else:
        delta = l_flag - mi
        mape = mape + delta
        mi = mape[a].min()
        mx = mape[a].max()
        mape = (mape - mi) / (mx - mi) * (h_flag - l_flag) + l_flag
    return mape

def mapping_diff(mape, mask_yes, mask_no, a):
    mask1 = a & mask_no # 车漆颜色 并且 不属于这个颜色
    mask2 = a & mask_yes # 车漆颜色 并且 属于这个颜色
    return 
    

def get_median(mape, mask):
    med = np.median(mape[mask])
    mask_1_4 = (mape < med) & mask
    mask_3_4 = (mape > med) & mask
    med_1_4 = np.median(mape[mask_1_4])
    med_3_4 = np.median(mape[mask_3_4])
    
    return med_1_4, med_3_4, med
    
    
# 汽车转色
def convert_color(hsv, mask, src_hsv, tag_hsv, tag_color, alpha, color_dict, mask_ori, pic_path, flag_s, flag_v, flag_h, light_conf, v_contrast, curve_conf):
    
    h, s, v = hsv[:, :, 0], hsv[:, :, 1], hsv[:, :, 2]
    tag_h, tag_s, tag_v = int(tag_hsv[0]), int(tag_hsv[1]), int(tag_hsv[2])
    
    a = mask
    b = mask == 0
    
    if 'white' in tag_color and flag_s ==  None:
        flag_s = [0, min(10, tag_s+1)]
        
    clahe = cv2.createCLAHE(clipLimit=v_contrast, tileGridSize=(8,8))
    
    delta_h = tag_hsv[0] - src_hsv[0]  
    delta_s = tag_hsv[1] - src_hsv[1]
    delta_v = tag_hsv[2] - src_hsv[2]
    print(f'delta_h: {delta_h}, delta_s: {delta_s}, delta_v: {delta_v}')


    h_ = h + delta_h
    s_ = s + delta_s 
    v_ = v + delta_v

    print('flag_s:{}, flag_v:{}'.format(flag_s, flag_v))
    # 映射到指定的区间
    h_ = mapping(h_, a, color_dict[tag_color][0][0], color_dict[tag_color][1][0], flag_h)
    s_ = mapping(s_, a, color_dict[tag_color][0][1], color_dict[tag_color][1][1], flag_s)
    v_ = mapping(v_, a, color_dict[tag_color][0][2], color_dict[tag_color][1][2], flag_v)

    # 变彩色 色调选填
    if judge_color(tag_color):  # 只调色调, s, v通过设置参数来调整
        if v_contrast != 0:
            v_1= clahe.apply((v_ * a).astype(np.uint8))
            hsv[:, :, 2] = v_1 * a + v * b
        else: 
            hsv[:, :, 2] = v_ * a + v * b
            
        hsv[:, :, 1] = s_ * a + s * b
        hsv[:, :, 0] = h_ * a + h * b
        res = save_result(hsv, alpha, mask_ori, pic_path, light_conf, curve_conf)
    # 变黑白灰
    else:
        # 黑白灰色调[0,180]不用管
        hsv[:, :, 0] = h_ * a + h * b
        hsv[:, :, 1] = s_ * a + s * b
        if v_contrast != 0:
            v_= clahe.apply((v_ * a).astype(np.uint8))

        hsv[:, :, 2] = v_ * a + v * b
        res = save_result(hsv, alpha, mask_ori, pic_path, light_conf, curve_conf)
    return res