import cv2
import numpy as np
import os

import os
import sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from src.utils import binarize
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit

# 定义输入图像和屏幕返回图像的文件夹路径
input_folder = 'imgs/grayscale_images_1'
returned_folder = 'imgs/returned_grayscales_1'

# 定义存储图像的文件名格式
# input_format = 'pat{}.bmp'
# output_format = '0_frame_{}.bmp'

# 存储输入灰度值和返回图像亮度的数组
input_gray_values = []
returned_brightness_values = []
mask = binarize(cv2.imread(os.path.join(input_folder, 'white.bmp'), cv2.IMREAD_GRAYSCALE))
# plt.imshow(mask, cmap='gray')
# 读取图像并计算灰度值和亮度值
error = 0
for i in range(256):
    # 生成文件名
    input_filename = os.path.join(input_folder, f'pat{i:03}.bmp')
    returned_filename = os.path.join(returned_folder, f'0_frame_{i}.bmp')
    
    # 读取输入灰度图和返回图像
    input_img = cv2.imread(input_filename, cv2.IMREAD_GRAYSCALE)
    returned_img = cv2.imread(returned_filename, cv2.IMREAD_GRAYSCALE)
    
    # 计算输入灰度值和返回图像的平均亮度值
    input_gray_value = np.mean(input_img)/255.0
    returned_brightness_value = np.mean(np.ma.masked_array(returned_img, mask=~mask))/255.0
    
    # 存储结果
    input_gray_values.append(input_gray_value)
    returned_brightness_values.append(returned_brightness_value)
    print(returned_brightness_value, input_gray_value)
    error += np.sum(np.abs(returned_brightness_value - input_gray_value))
print(error*255.0/256.0)

# 将结果转换为NumPy数组
input_gray_values = np.array(input_gray_values)
returned_brightness_values = np.array(returned_brightness_values)

# 定义拟合的Gamma函数
def gamma_function(x, gamma):
    return x ** gamma

# 使用scipy的curve_fit函数拟合Gamma值
popt, _ = curve_fit(gamma_function, input_gray_values, returned_brightness_values)
gamma_est = popt[0]

# 打印估计的Gamma值
print(f"Estimated Gamma: {gamma_est}")

# 绘制拟合曲线和数据点
plt.figure()
plt.scatter(input_gray_values, returned_brightness_values, label='Measured Data')
plt.plot(input_gray_values, gamma_function(input_gray_values, gamma_est), color='red', label=f'Fitted Gamma = {gamma_est:.2f}')
plt.xlabel('Input Gray Value')
plt.ylabel('Returned Brightness Value')
plt.title('Gamma Correction Fitting')
plt.legend()
plt.show()