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

def modify_b_channel_frequency(image, modify_pos, factor):
    # 提取B通道并进行FFT
    b_channel = image[:, :, 2]
    f_b = fftshift(fft2(b_channel))
    # 图像尺寸
    N, M = b_channel.shape
    amplitude = np.abs(f_b)
    phase = np.angle(f_b)
    # 获取修改位置的原始幅度
    x, y = modify_pos
    original_amplitude = np.abs(f_b[x, y])

    # 修改指定位置的频率幅度

    amplitude[x,y] = original_amplitude * factor
    modified_amplitude = amplitude[x, y]
    amplitude[N-x, M-y] = amplitude[N-x, M-y] * factor
    # modified_amplitude = original_amplitude
    f_b = amplitude * np.exp(1j * phase)

    # 执行逆FFT转换回空间域
    modified_b_channel = ifft2(ifftshift(f_b))
    modified_b_channel = np.real(modified_b_channel)


    # 保证修改后的B通道在合法范围内
    # modified_b_channel = np.clip(modified_b_channel, 0, 255).astype(np.uint8)

    # 合并修改后的B通道回原图
    # image[:, :, 2] = modified_b_channel

    # 返回修改前后的幅度信息
    return image, original_amplitude, modified_amplitude, modified_b_channel


def check_modification_effect(image, modify_pos):
    # 再次进行FFT以检查修改效果
    # b_channel = image[:, :, 2]
    b_channel = image
    f_b = fftshift(fft2(b_channel))
    x, y = modify_pos
    current_amplitude = np.abs(f_b[x, y])
    return current_amplitude, f_b


# 加载图像
img_path = 'demo_border.png'
image = np.array(Image.open(img_path))

# 设置修改位置和放大因子
modify_pos = (110, 125)  # 修改位置
factor = 2 # 放大原有的幅度

# 修改B通道并检查效果
modified_image, original_amplitude, modified_amplitude, modified_b_channel = modify_b_channel_frequency(image, modify_pos, factor)


image[:, :, 2] = modified_b_channel
filename = f"{datetime.now().strftime('%Y%m%d%H%M%S')}.png"
Image.fromarray(image).save(filename)

# final_image = Image.fromarray(image)
# final_image.save(filename)
#
img = Image.open(filename)
img = np.array(img)
b_channel = img[:, :, 2]

current_amplitude, f_b = check_modification_effect(modified_b_channel, modify_pos)


# 输出幅度对比
print(f"Original Amplitude: {original_amplitude}")
print(f"Modified Amplitude (should be): {modified_amplitude}")
print(f"Current Amplitude (after modifications and IFFT): {current_amplitude}")
# 打印出f_b中最大幅度
# print(f"Maximum Amplitude in f_b: {np.max(np.abs(f_b))}")

# 可视化显示
plt.figure(figsize=(12, 6))
plt.subplot(121)
plt.imshow(image)
plt.title('Original Image')
plt.subplot(122)
plt.imshow(modified_image)
plt.title('Modified Image')
plt.show()