import numpy as np
import hashlib
import time
from matplotlib.pyplot import imread, imsave


def generate_henon_sequence(size, a=1.4, b=0.3, x0=0.0, y0=0.0, discard=1000):
    """生成Henon Map混沌序列，用于置换索引。"""
    x, y = x0, y0
    sequence = []
    for _ in range(discard):  # 丢弃初始值以增加混沌
        x_new = 1 - a * x ** 2 + y
        y_new = b * x
        x, y = x_new, y_new
    for _ in range(size):
        x_new = 1 - a * x ** 2 + y
        y_new = b * x
        x, y = x_new, y_new
        sequence.append(x)
    return np.array(sequence)


def derive_henon_params(seed):
    np.random.seed(seed)
    a = 1.4 + np.random.uniform(-0.01, 0.01)   # 更小的扰动
    b = 0.3 + np.random.uniform(-0.005, 0.005) # 更小的扰动
    x0 = np.random.uniform(-0.1, 0.1)           # 限制初始值范围
    y0 = np.random.uniform(-0.1, 0.1)
    return a, b, x0, y0

def obfuscate(input_path, output_path, iterations=1, seed=None, password=None):
    img = imread(input_path)
    if len(img.shape) == 2:
        img = img[..., np.newaxis]  # 灰度转3D
    H, W, C = img.shape

    if password is not None:
        seed = int(hashlib.sha256(password.encode()).hexdigest(), 16) % (2 ** 32)

    total_pixels = H * W
    for it in range(iterations):
        if seed is not None:
            effective_seed = (seed + it) % (2 ** 32)
            a, b, x0, y0 = derive_henon_params(effective_seed)
        else:
            a, b, x0, y0 = 1.4, 0.3, 0.0, 0.0
        seq = generate_henon_sequence(total_pixels, a, b, x0, y0)
        indices = np.argsort(seq)  # 使用排序后的序列作为置换
        flat_img = img.reshape((total_pixels, C))
        shuffled = flat_img[indices]
        img = shuffled.reshape((H, W, C))

    if C == 1:
        img = img[:, :, 0]
    print(f"HenonMap 混淆完成，已保存到：{output_path}")
    imsave(output_path, img)


def deobfuscate(input_path, output_path, iterations=1, seed=None, password=None):
    img = imread(input_path)
    if len(img.shape) == 2:
        img = img[..., np.newaxis]
    H, W, C = img.shape

    if password is not None:
        seed = int(hashlib.sha256(password.encode()).hexdigest(), 16) % (2 ** 32)

    total_pixels = H * W
    for it in range(iterations - 1, -1, -1):  # 逆序恢复
        if seed is not None:
            effective_seed = (seed + it) % (2 ** 32)
            a, b, x0, y0 = derive_henon_params(effective_seed)
        else:
            a, b, x0, y0 = 1.4, 0.3, 0.0, 0.0
        seq = generate_henon_sequence(total_pixels, a, b, x0, y0)
        indices = np.argsort(seq)
        # 创建逆置换
        inv_indices = np.zeros(total_pixels, dtype=int)
        inv_indices[indices] = np.arange(total_pixels)
        flat_img = img.reshape((total_pixels, C))
        restored = flat_img[inv_indices]
        img = restored.reshape((H, W, C))

    if C == 1:
        img = img[:, :, 0]
    print(f"HenonMap 解析完成，已保存到：{output_path}")
    imsave(output_path, img)