import os
import cv2
import numpy as np

# 高斯模糊
def image_blur(image,ksize):
    if image is None:
        raise ValueError("Image is not loaded correctly.")
    gaussian_blur = cv2.GaussianBlur(image, (ksize, ksize), 0)
    return gaussian_blur

# 椒盐噪声
def add_salt_pepper(image, salt_prob=0.05, pepper_prob=0.05):
    noisy = image.copy()
    # 添加盐噪声
    salt = np.random.randint(0, 256, size=image.shape[:2])
    noisy[salt < salt_prob*255] = 255
    # 添加椒噪声
    pepper = np.random.randint(0, 256, size=image.shape[:2])
    noisy[pepper < pepper_prob*255] = 0
    return noisy

# Brenner 梯度函数计算
def brenner(img):
    '''
    :param img: narray 二维灰度图像
    :return: int 图像越清晰越大
    '''
    return int(np.sum((img[:-2, :] - img[2:, :]) ** 2))

# Laplacian 梯度函数计算
def Laplacian(img):
    '''
    :param img: narray 二维灰度图像
    :return: int 图像越清晰越大
    '''
    return int(cv2.Laplacian(img, cv2.CV_64F).var())

# SMD 梯度函数计算
def SMD(img):
    '''
    :param img: narray 二维灰度图像
    :return: int 图像越清晰越大
    '''
    diff_x = np.abs(img[1:-1, :] - img[:-2, :])
    diff_y = np.abs(img[1:, :] - img[:-1, :])
    return int(np.sum(diff_x) + np.sum(diff_y))

# SMD2 函数计算
def SMD2(img):
    '''
    :param img: narray 二维灰度图像
    :return: int 图像越清晰越大
    '''
    sobel_x = cv2.Sobel(img, cv2.CV_64F, 1, 0, ksize=3)
    sobel_y = cv2.Sobel(img, cv2.CV_64F, 0, 1, ksize=3)
    gradient = np.abs(sobel_x) + np.abs(sobel_y)
    return int(np.sum(gradient ** 2))

# 方差函数计算
def variance(img):
    '''
    :param img: narray 二维灰度图像
    :return: int 图像越清晰越大
    '''
    u = np.mean(img)
    return int(np.sum((img - u) ** 2))

# Energy 函数计算
def energy(img):
    '''
    :param img: narray 二维灰度图像
    :return: int 图像越清晰越大
    '''
    diff_x = img[1:, :-1] - img[:-1, :-1]
    diff_y = img[:-1, 1:] - img[:-1, :-1]
    return int(np.sum(diff_x ** 2 * diff_y ** 2))

# Vollath 函数计算
def Vollath(img):
    '''
    :param img: narray 二维灰度图像
    :return: int 图像越清晰越大
    '''
    u = np.mean(img)
    return int(-img.size * (u ** 2) + np.sum(img[:-1, :] * img[1:, :]))

# Entropy 函数计算
def entropy(img):
    '''
    :param img: narray 二维灰度图像
    :return: int 图像越清晰越大
    '''
    hist = np.bincount(img.flatten(), minlength=256)
    p = hist / img.size
    p = p[p > 0]
    return int(-np.sum(p * np.log(p)))



def image_test(filename,image,noisy_img,blur_image):
    image_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    noisy_gray = cv2.cvtColor(noisy_img, cv2.COLOR_BGR2GRAY)
    blur_gray = cv2.cvtColor(blur_image, cv2.COLOR_BGR2GRAY)

    # 定义一个格式化字符串，确保每个数值占据相同的宽度
    format_str = "{:<10} {:<15} {:<15} {:<15}"

    # 打印标题
    print(format_str.format(filename, "Original", "Noisy", "Blur"))

    # 打印每个指标的值
    print(format_str.format("Brenner", brenner(image_gray), brenner(noisy_gray), brenner(blur_gray)))
    print(format_str.format("Laplacian", Laplacian(image_gray), Laplacian(noisy_gray), Laplacian(blur_gray)))
    print(format_str.format("SMD", SMD(image_gray), SMD(noisy_gray), SMD(blur_gray)))
    print(format_str.format("SMD2", SMD2(image_gray), SMD2(noisy_gray), SMD2(blur_gray)))
    print(format_str.format("Variance", variance(image_gray), variance(noisy_gray), variance(blur_gray)))
    print(format_str.format("Energy", energy(image_gray), energy(noisy_gray), energy(blur_gray)))
    print(format_str.format("Vollath", Vollath(image_gray), Vollath(noisy_gray), Vollath(blur_gray)))
    print(format_str.format("Entropy", entropy(image_gray), entropy(noisy_gray), entropy(blur_gray)))


    # print(filename,'noisy','blur')
    # print('Brenner',brenner(image_gray),brenner(noisy_gray),brenner(blur_gray))
    # print('Laplacian',Laplacian(image_gray),Laplacian(noisy_gray),Laplacian(blur_gray))
    # print('SMD',SMD(image_gray), SMD(noisy_gray), SMD(blur_gray))
    # print('SMD2',SMD2(image_gray), SMD2(noisy_gray), SMD2(blur_gray))
    # print('Variance',variance(image_gray),variance(noisy_gray),variance(blur_gray))
    # print('Energy',energy(image_gray),energy(noisy_gray),energy(blur_gray))
    # print('Vollath',Vollath(image_gray),Vollath(noisy_gray),Vollath(blur_gray))
    # print('Entropy',entropy(image_gray),entropy(noisy_gray),entropy(blur_gray))
  


def process_images(input_dir, salty_dir,blur_dir):
    # 创建输出目录（如果不存在）
    if not os.path.exists(salty_dir):
        os.makedirs(salty_dir)
    if not os.path.exists(blur_dir):
        os.makedirs(blur_dir)
    
    # 支持的图片格式
    img_exts = ('.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff')
    
    # 遍历输入目录
    for filename in os.listdir(input_dir):
        if filename.lower().endswith(img_exts):
            # 读取图片
            img_path = os.path.join(input_dir, filename)
            img = cv2.imread(img_path)
            img = cv2.resize(img,(512, 512))

            if img is None:
                raise ValueError("Error: Image not found or unable to load.")
            
            if img is not None:

                # 添加椒盐噪声
                noisy_img = add_salt_pepper(img)
                # 添加Gauss噪声
                blur_image = image_blur(img,ksize=25)



                # 构造新文件名
                name, ext = os.path.splitext(filename)
                salty_filename = f"{name}_salty{ext}"
                blur_filename = f"{name}_blur{ext}"

                image_test(filename,img,noisy_img,blur_image)
                print('\n')


                salty_path = os.path.join(salty_dir, salty_filename)
                cv2.imwrite(salty_path, noisy_img)
                blur_path = os.path.join(blur_dir, blur_filename)
                cv2.imwrite(blur_path, blur_image)


if __name__ == '__main__':
    input_dir = 'data_process/Clarity Test/images/source'
    salty_dir = 'data_process/Clarity Test/images/salty'
    blur_dir = 'data_process/Clarity Test/images/blur'
    process_images(input_dir, salty_dir,blur_dir)