import os
import cv2
import numpy as np
from image_module import get_lbp_entropy, get_dct_high_freq_energy, get_fft_symmetry_error


def process_folder(folder_path, label):
    results = []
    print(f"正在分析 [{label}] 样本集: {folder_path}")

    for filename in os.listdir(folder_path):
        if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
            filepath = os.path.join(folder_path, filename)
            # 读取图片并转换为灰度图
            img = cv2.imdecode(np.fromfile(filepath, dtype=np.uint8), cv2.IMREAD_COLOR)
            if img is None: continue
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

            # 提取特征
            lbp = get_lbp_entropy(gray)
            dct = get_dct_high_freq_energy(gray)
            fft = get_fft_symmetry_error(gray)

            results.append((lbp, dct, fft))

    if not results:
        return None

    results_arr = np.array(results)
    print(f"--- {label} 统计结果 ({len(results)}张) ---")
    print(
        f"LBP 熵  : 平均 {np.mean(results_arr[:, 0]):.3f} | 范围 [{np.min(results_arr[:, 0]):.3f} - {np.max(results_arr[:, 0]):.3f}]")
    print(
        f"DCT 占比: 平均 {np.mean(results_arr[:, 1]):.3f} | 范围 [{np.min(results_arr[:, 1]):.3f} - {np.max(results_arr[:, 1]):.3f}]")
    print(
        f"FFT 误差: 平均 {np.mean(results_arr[:, 2]):.3f} | 范围 [{np.min(results_arr[:, 2]):.3f} - {np.max(results_arr[:, 2]):.3f}]\n")
    return results_arr


if __name__ == "__main__":
    # 请在当前目录下新建这两个文件夹，分别放几十张真实的和AI生成的图进去
    real_folder = "./data/real"
    ai_folder = "./data/ai"

    print("开始执行特征分布寻优...\n")
    process_folder(real_folder, "真实图像")
    process_folder(ai_folder, "AI生成图像")