import argparse
import cv2
import numpy as np

def find_last_x(image, slope_threshold = 1000):
    x,y = [],[]
    hist, bins = np.histogram(image, bins=256, range=[0, 256])

    #找到50以内的最高峰
    max_y , max_i = 0, 5
    for i in range(5, 25):
        if hist[i] > max_y:
            max_y = hist[i]
            max_i = i
    print(f'50以内最高峰值y：{max_y}，最高峰位置x：{max_i}')

    for i in range(2, max_i):
        x.append(i)
        y.append(hist[i])
    slopes = [abs(y[i + 1] - y[i]) for i in range(len(x) - 1)]

    current_interval = []
    max_interval = []
    max_x = {}
    for i, slope in enumerate(slopes):
        current_interval.append(slope)
        if slope >= slope_threshold:
            if len(current_interval) > len(max_interval):
                max_interval = current_interval.copy()
                max_x[x[i]] = slope
            current_interval = []

    print(max_x)
    last_x = list(max_x)[-1]
    last_y = max_x[last_x]
    return last_x, last_y


def find_last_high(image, slope_threshold = 2500):
    x = []
    y = []
    hist, bins = np.histogram(image, bins=255, range=[2, 255])

    #找到200以上的最高峰
    max_y = 0
    max_i = 254
    for i in range(220, 255):
        if hist[i] > max_y:
            max_y = hist[i]
            max_i = i
    print(f'200以上的最高峰值y：{max_y}，最高峰位置x：{max_i}')

    for i in range(max_i, 255):
        x.append(i)
        y.append(hist[i])
    slopes = [abs(y[i + 1] - y[i]) for i in range(len(x) - 1)]

    current_interval = []
    max_interval = []
    max_x = {}
    find = False
    for i in range(len(slopes) - 1, -1, -1):
        slope = slopes[i]
        current_interval.append(slope)
        if slope >= slope_threshold:
            find = True
            if len(current_interval) > len(max_interval):
                max_interval = current_interval.copy()
                max_x[x[i]] = slope
            current_interval = []
        #如果没有找到200以上很平，而且高度小于5000，就按220位置削平
        if not find and hist[220] < 5000:
            max_x[220] = hist[220]

    print(max_x)
    if len(max_x) > 0:
        last_x = list(max_x)[0]
        last_y = max_x[last_x]
    else:
        print(f'找不到200以上曲线较平的区间，使用254作为最高峰')
        last_x = 254
        last_y = hist[254]
    return last_x, last_y


def ps_color_scale_adjustment(image, shadow=0, highlight=255, midtones=1):
    '''
    模拟 PS 的色阶调整； 0 <= Shadow < Highlight <= 255
    :param image: 传入的图片
    :param shadow: 黑场(0-Highlight)
    :param highlight: 白场(Shadow-255)
    :param midtones: 灰场(9.99-0.01)
    :return: 图片
    '''
    if highlight > 255:
        highlight = 255
    if shadow < 0:
        shadow = 0
    if shadow >= highlight:
        shadow = highlight - 2
    if midtones > 9.99:
        midtones = 9.99
    if midtones < 0.01:
        midtones = 0.01
    image = np.array(image, dtype=np.float16)
    # 计算白场 黑场离差
    Diff = highlight - shadow
    image = image - shadow
    image[image < 0] = 0
    image = (image / Diff) ** (1 / midtones) * 255
    image[image > 255] = 255
    image = np.array(image, dtype=np.uint8)

    return image


def white_purification_utils(image_path):
    input_image = cv2.imread(image_path)
    # low_x_thresh, low_y_frequency = low_find_histogram_range(input_image, low_y_limit)
    low_x_thresh, low_y_frequency = find_last_x(input_image)
    # high_x_thresh, high_y_frequency = high_find_histogram_range(input_image, high_y_limit)
    high_x_thresh, high_y_frequency = find_last_high(input_image)
    print(f"{low_x_thresh} 区间, {low_y_frequency} 频次")
    print(f"{high_x_thresh} 区间, {high_y_frequency} 频次")
    high_output_image = ps_color_scale_adjustment(input_image, shadow=low_x_thresh, highlight=high_x_thresh, midtones=1)
    file_extension = image_path.lower().split('.')[-1]
    if file_extension == 'png':
        cv2.imwrite(image_path, high_output_image)  # PNG格式无需压缩参数
    else:
        # cv2.imwrite(image_path.replace(".jpg", '_white.jpg'), high_output_image, [cv2.IMWRITE_JPEG_QUALITY, 95])  # 保存图片的质量是原图的 95%
        cv2.imwrite(image_path, high_output_image, [cv2.IMWRITE_JPEG_QUALITY, 95])  # 保存图片的质量是原图的 95%


# if __name__ == "__main__":
#     parser = argparse.ArgumentParser()
#     parser.add_argument("-i","--image_path", type=str, default="")
#     args = parser.parse_args()
#     white_purification_utils(args.image_path)