import cv2
import numpy as np
import matplotlib.pyplot as plt
from scipy import ndimage
from scipy.signal import convolve2d

def wiener_filter_builtin(image_path, blur_kernel_size=5, noise_power=0.01):
    """
    使用OpenCV实现维纳滤波（通过去卷积近似）
    :param image_path: 输入图像路径
    :param blur_kernel_size: 模糊核大小
    :param noise_power: 噪声功率
    :return: 原图和滤波后的图像
    """
    # 读取图像
    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE).astype(np.float32)
    
    # 创建模糊核（模拟退化过程）
    kernel = np.ones((blur_kernel_size, blur_kernel_size), np.float32) / (blur_kernel_size ** 2)
    
    # 应用维纳滤波
    # OpenCV没有直接的维纳滤波函数，这里使用去卷积近似
    deconvolved = cv2.deconvblind(img, kernel)[0]
    
    return cv2.cvtColor(img.astype(np.uint8), cv2.COLOR_GRAY2BGR), cv2.cvtColor(deconvolved.astype(np.uint8), cv2.COLOR_GRAY2BGR)

def manual_wiener_filter(image_path, kernel_size=5, noise_var=0.01):
    """
    手动实现维纳滤波
    :param image_path: 输入图像路径
    :param kernel_size: 模糊核大小
    :param noise_var: 噪声方差
    :return: 维纳滤波后的图像
    """
    # 读取图像
    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE).astype(np.float32)
    height, width = img.shape
    
    # 创建模糊核（模拟退化过程）
    kernel = np.ones((kernel_size, kernel_size), np.float32) / (kernel_size ** 2)
    
    # 对图像和核进行零填充，使其大小相同
    padded_kernel = np.zeros_like(img)
    kh, kw = kernel.shape
    padded_kernel[:kh, :kw] = kernel
    
    # 将核移到中心位置（为了FFT）
    padded_kernel = np.roll(padded_kernel, -kh//2, axis=0)
    padded_kernel = np.roll(padded_kernel, -kw//2, axis=1)
    
    # 计算FFT
    img_fft = np.fft.fft2(img)
    kernel_fft = np.fft.fft2(padded_kernel)
    
    # 计算维纳滤波器
    kernel_power = np.abs(kernel_fft) ** 2
    wiener_filter = np.conj(kernel_fft) / (kernel_power + noise_var)
    
    # 应用维纳滤波器
    filtered_fft = img_fft * wiener_filter
    
    # 逆FFT
    filtered = np.real(np.fft.ifft2(filtered_fft))
    
    # 确保像素值在有效范围内
    filtered = np.clip(filtered, 0, 255)
    
    return cv2.cvtColor(img.astype(np.uint8), cv2.COLOR_GRAY2BGR), cv2.cvtColor(filtered.astype(np.uint8), cv2.COLOR_GRAY2BGR)

def motion_deblur_wiener(image_path, motion_angle=0, motion_dist=10):
    """
    运动模糊的维纳滤波去模糊
    :param image_path: 输入图像路径
    :param motion_angle: 运动角度（度）
    :param motion_dist: 运动距离
    :return: 原图和去模糊后的图像
    """
    # 读取图像
    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE).astype(np.float32)
    
    # 创建运动模糊核
    def motion_kernel(size, angle, dist):
        kernel = np.zeros((size, size))
        center = size // 2
        
        # 计算运动方向的单位向量
        rad = np.deg2rad(angle)
        dx = np.cos(rad)
        dy = np.sin(rad)
        
        # 在运动路径上设置值
        for i in range(dist):
            x = int(center + dx * i)
            y = int(center + dy * i)
            if 0 <= x < size and 0 <= y < size:
                kernel[y, x] = 1
        
        # 归一化
        kernel = kernel / np.sum(kernel)
        return kernel
    
    # 生成运动模糊核
    kernel_size = max(15, motion_dist + 5)  # 确保核足够大
    motion_kernel_mat = motion_kernel(kernel_size, motion_angle, motion_dist)
    
    # 对图像和核进行零填充
    padded_kernel = np.zeros_like(img)
    kh, kw = motion_kernel_mat.shape
    start_row = max(0, (height - kh) // 2)
    start_col = max(0, (width - kw) // 2)
    padded_kernel[start_row:start_row+kh, start_col:start_col+kw] = motion_kernel_mat
    
    # 计算FFT
    img_fft = np.fft.fft2(img)
    kernel_fft = np.fft.fft2(padded_kernel)
    
    # 维纳滤波参数
    noise_power = 0.01
    kernel_power = np.abs(kernel_fft) ** 2
    
    # 避免除零
    wiener_filter = np.zeros_like(kernel_fft)
    mask = kernel_power > 1e-8
    wiener_filter[mask] = np.conj(kernel_fft[mask]) / (kernel_power[mask] + noise_power)
    
    # 应用维纳滤波器
    filtered_fft = img_fft * wiener_filter
    
    # 逆FFT
    filtered = np.real(np.fft.ifft2(filtered_fft))
    
    # 确保像素值在有效范围内
    filtered = np.clip(filtered, 0, 255)
    
    return cv2.cvtColor(img.astype(np.uint8), cv2.COLOR_GRAY2BGR), cv2.cvtColor(filtered.astype(np.uint8), cv2.COLOR_GRAY2BGR)

def gaussian_deblur_wiener(image_path, sigma=2):
    """
    高斯模糊的维纳滤波去模糊
    :param image_path: 输入图像路径
    :param sigma: 高斯核的标准差
    :return: 原图和去模糊后的图像
    """
    # 读取图像
    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE).astype(np.float32)
    height, width = img.shape
    
    # 创建高斯模糊核
    def gaussian_kernel(size, sigma):
        kernel = np.zeros((size, size))
        center = size // 2
        for i in range(size):
            for j in range(size):
                x, y = i - center, j - center
                kernel[i, j] = np.exp(-(x**2 + y**2) / (2 * sigma**2))
        return kernel / np.sum(kernel)
    
    kernel_size = int(2 * np.ceil(3 * sigma) + 1)  # 确保核大小合适
    gaussian_kernel_mat = gaussian_kernel(kernel_size, sigma)
    
    # 对图像和核进行零填充
    padded_kernel = np.zeros_like(img)
    kh, kw = gaussian_kernel_mat.shape
    start_row = max(0, (height - kh) // 2)
    start_col = max(0, (width - kw) // 2)
    padded_kernel[start_row:start_row+kh, start_col:start_col+kw] = gaussian_kernel_mat
    
    # 计算FFT
    img_fft = np.fft.fft2(img)
    kernel_fft = np.fft.fft2(padded_kernel)
    
    # 维纳滤波参数
    noise_power = 0.001
    kernel_power = np.abs(kernel_fft) ** 2
    
    # 避免除零
    wiener_filter = np.zeros_like(kernel_fft)
    mask = kernel_power > 1e-8
    wiener_filter[mask] = np.conj(kernel_fft[mask]) / (kernel_power[mask] + noise_power)
    
    # 应用维纳滤波器
    filtered_fft = img_fft * wiener_filter
    
    # 逆FFT
    filtered = np.real(np.fft.ifft2(filtered_fft))
    
    # 确保像素值在有效范围内
    filtered = np.clip(filtered, 0, 255)
    
    return cv2.cvtColor(img.astype(np.uint8), cv2.COLOR_GRAY2BGR), cv2.cvtColor(filtered.astype(np.uint8), cv2.COLOR_GRAY2BGR)

def add_motion_blur(image, angle=0, distance=10):
    """
    为图像添加运动模糊
    :param image: 输入图像
    :param angle: 运动角度
    :param distance: 运动距离
    :return: 添加运动模糊后的图像
    """
    # 创建运动模糊核
    kernel = np.zeros((distance, distance))
    center = distance // 2
    
    # 计算运动方向的单位向量
    rad = np.deg2rad(angle)
    dx = np.cos(rad)
    dy = np.sin(rad)
    
    # 在运动路径上设置值
    for i in range(distance):
        x = int(center + dx * i)
        y = int(center + dy * i)
        if 0 <= x < distance and 0 <= y < distance:
            kernel[y, x] = 1
    
    # 归一化
    kernel = kernel / np.sum(kernel)
    
    # 应用卷积
    blurred = cv2.filter2D(image, -1, kernel)
    
    return blurred

# 使用示例
if __name__ == "__main__":
    # 注意：需要提供实际的图像路径
    # img, result = wiener_filter_builtin('image.jpg', 5, 0.01)
    # manual_img, manual_result = manual_wiener_filter('image.jpg', 5, 0.01)
    # motion_img, motion_result = motion_deblur_wiener('image.jpg', 45, 10)
    # gaussian_img, gaussian_result = gaussian_deblur_wiener('image.jpg', 2)
    pass