import numpy as np
from numpy.fft import fft2, ifft2, fftshift, ifftshift
import matplotlib.pyplot as plt
from PIL import Image


def modify_amplitude_only(image_path, u, v, factor):
    """
    修改图像B通道在(u, v)位置的频率分量的幅度，并保持相位不变。

    参数:
    image_path: 输入图像的路径。
    u, v: 要修改的频率分量的位置。
    factor: 幅度修改因子。

    返回:
    original_amplitude: 原始幅度。
    modified_amplitude: 修改后的预期幅度。
    current_amplitude: 实际当前幅度。
    """
    # 读取图像
    image = np.array(Image.open(image_path))
    B_channel = image[:, :, 2]  # 提取B通道

    N, M = B_channel.shape
    F = fftshift(fft2(B_channel))

    # 获取原始的幅度和相位
    original_amplitude = np.abs(F[u, v])

    # 修改幅度
    amplitude = np.abs(F)
    phase = np.angle(F)
    amplitude[u, v] *= factor
    amplitude[N - u, M - v] *= factor  # 确保共轭对称
    modified_amplitude = amplitude[u, v]
    # 重新构造频域数据
    F_new = amplitude * np.exp(1j * phase)
    modified_B_channel = np.real(ifft2(ifftshift(F_new)))  # 执行逆FFT并取实部

    # 合并修改后的B通道回原图
    modified_image = image.copy()
    modified_image[:, :, 2] = modified_B_channel.astype(np.float64)

    # 从修改后的完整图像中提取B通道并进行DFT
    modified_B_channel_from_image = modified_image[:, :, 2]
    F_check = fftshift(fft2(modified_B_channel_from_image))
    current_amplitude = np.abs(F_check[u, v])

    return original_amplitude, modified_amplitude, current_amplitude


# 示例使用
image_path = 'demo_border.png'  # 你的图像路径
u, v = 128, 128  # 选择要修改的频点位置
factor = 20  # 幅度增加因子

# 修改图像并获取数据
original_amplitude, modified_amplitude, current_amplitude = modify_amplitude_only(image_path, u, v, factor)

# 输出幅度数据对比
print(f"Original Amplitude at ({u}, {v}): {original_amplitude}")
print(f"Modified Amplitude (expected) at ({u}, {v}): {modified_amplitude}")
print(f"Current Amplitude (after modification and IFFT) at ({u}, {v}): {current_amplitude}")