import numpy as np
from scipy import fftpack

def robust_watermark_encode(payload: str, key: bytes, img_shape: tuple) -> np.ndarray:
    """
    生成抗攻击水印模板（频域扩频+加密）
    :param payload: 版权ID/时间戳等结构化数据
    :param key: 16字节AES密钥（保障水印唯一性）
    :param img_shape: 目标图像尺寸 (H, W)
    :return: 归一化水印模板 (H, W)
    """
    # 1. 信息编码：结构化数据→二进制流
    binary_stream = ''.join(format(b, '08b') for b in payload.encode('utf-8'))
    bits = np.array([int(b) for b in binary_stream])

    # 2. 扩频处理：每位扩展为伪随机序列（抗裁剪/噪声）
    np.random.seed(int.from_bytes(key[:4], 'big'))
    spreading_code = np.random.choice([-1, 1], size=(len(bits), 100))
    spread_seq = np.array([code if bit else -code for bit, code in zip(bits, spreading_code)]).flatten()

    # 3. 频域映射：生成与图像尺寸匹配的水印模板
    template = np.zeros(img_shape)
    # 将扩频序列映射至DCT中高频区域（人眼不敏感区）
    h, w = img_shape
    dct_block = fftpack.dct(template, axis=0, norm='ortho')
    dct_block[h//4:h//2, w//4:w//2] = spread_seq[: (h//4)*(w//4)].reshape(h//4, w//4) * 0.05  # 微扰强度

    # 4. 逆变换回空间域（作为嵌入引导信号）
    watermark_template = fftpack.idct(dct_block, axis=0, norm='ortho')
    return watermark_template / np.max(np.abs(watermark_template))  # 归一化

def watermark_guided_denoising(
    unet, scheduler, latents, watermark_template, 
    guidance_scale=0.3, timesteps=None
):
    """
    在DDIM采样中动态嵌入水印（修改噪声预测梯度）
    :param watermark_template: 步骤一生成的归一化模板
    :param guidance_scale: 水印强度（0.1~0.5，需实验调优）
    """
    for t in timesteps:
        # 标准噪声预测
        noise_pred = unet(latents, t).sample

        # 水印引导：计算潜变量空间的梯度扰动
        # 将水印模板映射至当前潜变量尺度
        wm_latent = F.interpolate(
            torch.tensor(watermark_template)[None, None], 
            size=latents.shape[-2:], mode='bilinear'
        ).to(latents.device)

        # 关键：沿水印梯度方向微调噪声预测（非简单加法！）
        perturbation = guidance_scale * wm_latent * torch.sign(noise_pred)
        noise_pred = noise_pred + perturbation

        # 标准采样更新
        latents = scheduler.step(noise_pred, t, latents).prev_sample

    return latents

def forensic_watermark_extract(image: np.ndarray, original_template: np.ndarray) -> dict:
    """
    司法级水印验证：含攻击模拟与置信度评估
    """
    # 1. 预处理：统一尺寸、色彩空间
    img_proc = preprocess_for_extraction(image)

    # 2. 攻击模拟（提升提取鲁棒性）：
    #    - 添加轻微高斯噪声（模拟传输损失）
    #    - 轻度JPEG压缩（QF=90）再解压
    img_sim = simulate_common_attacks(img_proc)

    # 3. 频域相关性检测
    img_dct = fftpack.dct(img_sim, axis=0, norm='ortho')
    wm_region = img_dct[h//4:h//2, w//4:w//2].flatten()

    # 4. 计算与原始模板的相关系数（核心指标）
    correlation = np.corrcoef(wm_region, original_template.flatten())[0, 1]

    # 5. 置信度分级（司法场景关键）
    if correlation > 0.75:
        confidence = "高置信（可作为证据）"
    elif correlation > 0.6:
        confidence = "中置信（需辅助验证）"
    else:
        confidence = "未检测到有效水印"

    return {
        "correlation": round(correlation, 4),
        "confidence_level": confidence,
        "timestamp": datetime.utcnow().isoformat()
    }