import matplotlib
matplotlib.use('Agg')  # 使用非 GUI 后端

import cv2
import numpy as np
import matplotlib.pyplot as plt

def analyze_surface_roughness(image_path):
    # 读取图像
    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    if img is None:
        raise FileNotFoundError(f"无法加载图像: {image_path}")

    # 预处理
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    img = clahe.apply(img)
    img = cv2.medianBlur(img, 5)

    # FFT
    f = np.fft.fft2(img)
    fshift = np.fft.fftshift(f)
    magnitude_spectrum = 20 * np.log(np.abs(fshift) + 1e-9)

    # 动态低频区域
    rows, cols = img.shape
    crow, ccol = rows // 2, cols // 2
    radius = min(rows, cols) // 20
    mask = np.ones((rows, cols), np.uint8)
    cv2.circle(mask, (ccol, crow), radius, 0, thickness=-1)

    # 计算能量占比
    total_energy = np.sum(magnitude_spectrum)
    high_freq_energy = np.sum(magnitude_spectrum * mask)
    ratio = high_freq_energy / total_energy

    # 可视化并保存
    plt.figure(figsize=(12, 6))

    plt.subplot(1, 3, 1)
    plt.imshow(img, cmap='gray')
    plt.title('预处理后图像')
    plt.axis('off')

    plt.subplot(1, 3, 2)
    plt.imshow(magnitude_spectrum, cmap='gray')
    plt.title('频谱图（中心为低频）')
    plt.axis('off')

    overlay = magnitude_spectrum.copy()
    if len(overlay.shape) == 2:
        overlay = np.stack([overlay]*3, axis=-1)
    cv2.circle(overlay, (ccol, crow), radius, (1, 0, 0), 2)
    plt.subplot(1, 3, 3)
    plt.imshow(overlay)
    plt.title(f'低频区域（红圈内）\n高频能量占比: {ratio:.3f}')
    plt.axis('off')

    plt.tight_layout()
    plt.savefig("fft_analysis_result.png", dpi=150, bbox_inches='tight')
    plt.close()  # 释放内存

    # 输出结果
    print(f"✅ 图像尺寸: {cols}x{rows}")
    print(f"📊 高频能量占比: {ratio:.3f}")
    if ratio < 0.8:
        print("🟢 评价: 表面较光滑")
    elif ratio < 0.95:
        print("🟡 评价: 表面中等粗糙")
    else:
        print("🔴 评价: 表面非常粗糙")

    return ratio

# 运行
if __name__ == "__main__":
    analyze_surface_roughness("1.jpg")