# This is a sample Python script.

# Press ⌃R to execute it or replace it with your code.
# Press Double ⇧ to search everywhere for classes, files, tool windows, actions, and settings.
import numpy as np
from scipy.fft import fft2, fftshift
import cv2
import matplotlib.pyplot as plt

def print_hi(name):
    # Use a breakpoint in the code line below to debug your script.
    print(f'Hi, {name}')  # Press ⌘F8 to toggle the breakpoint.


def encode_watermark_to_binary(watermark_sequence):
    binary_sequence = []
    for number in watermark_sequence:
        # 初始化一个100位全为'0'的二进制字符串
        binary_100_bit = ['0'] * 100
        # 在数字对应的位置设置'1'
        if 0 <= number < 100:
            binary_100_bit[number-1] = '1'
        # 连接这些位，形成一个100位的二进制字符串
        binary_sequence.append(''.join(binary_100_bit))
    return binary_sequence

def process_image_b_component(image_path):
    """
    处理图像的B通道，执行DFT并移动频谱。

    :param image_path: 输入图像的路径。
    :return: 移位后DFT的幅度矩阵。
    """
    # 使用合适的库加载图像（例如，PIL, OpenCV）
    # 这里我们会加载图像并提取B通道
    # 作为演示，我们将创建一个假的256x256 B通道矩阵

    #根据image_path加载图像
    image = cv2.imread(image_path)
    # 检查图像是否成功加载
    if image is None:
        raise ValueError("无法加载图像，请检查路径是否正确。")

    # 确保图像的大小是256x256
    image_resized = cv2.resize(image, (256, 256))

    # 提取B通道
    b_channel = image_resized[:, :, 0]

    # 对B通道矩阵执行DFT
    dft_result = fft2(b_channel)

    # 将零频率分量移动到频谱的中心
    dft_shifted = fftshift(dft_result)

    # 获取移位后DFT的幅度
    magnitude = np.abs(dft_shifted)

    return magnitude

def decode_binary_to_watermark(binary_sequence):
    """
    Decode the binary sequence back into the original array of numbers.

    :param binary_sequence: A list of 100-bit strings representing the binary codes.
    :return: The original array of numbers.
    """
    original_numbers = []
    for binary_code in binary_sequence:
        if '1' in binary_code:
            # Find the index of '1' and convert it back to the original number
            position = binary_code.index('1') + 1  # Add 1 because the index starts at 0
            original_number = position
            original_numbers.append(original_number)
    return original_numbers

# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    # print_hi('PyCharm')
    watermark_sequence = [12, 34, 56, 44, 90, 55, 34, 56]
    result = encode_watermark_to_binary(watermark_sequence)
    print(result)
    # # 打印编码后的二进制序列
    # for binary in result:
    #     print(binary)

    watermark_sequence = decode_binary_to_watermark(result)
    print(watermark_sequence)

# See PyCharm help at https://www.jetbrains.com/help/pycharm/