#!/usr/bin/env python3
"""
測試超像素區域保存功能
Test script for superpixel region saving functionality
"""

import os
import sys
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from skimage.segmentation import slic
from skimage.util import img_as_float

# Add current directory to path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from Density_Peak_Algorithm import save_superpixel_regions, compute_decision_graph

def test_superpixel_regions_save():
    """測試超像素區域保存功能"""
    
    print("🧪 開始測試超像素區域保存功能...")
    
    # 創建測試目錄
    test_dir = "test_superpixel_regions_results"
    os.makedirs(test_dir, exist_ok=True)
    
    # 創建一個簡單的測試圖像
    test_image_path = os.path.join(test_dir, "test_image.png")
    
    # 創建一個彩色測試圖像 (200x200)
    height, width = 200, 200
    test_img = np.zeros((height, width, 3), dtype=np.uint8)
    
    # 添加不同顏色的區域
    test_img[:100, :100] = [255, 0, 0]    # 紅色
    test_img[:100, 100:] = [0, 255, 0]    # 綠色
    test_img[100:, :100] = [0, 0, 255]    # 藍色
    test_img[100:, 100:] = [255, 255, 0]  # 黃色
    
    # 添加一些噪聲
    noise = np.random.randint(-30, 30, (height, width, 3))
    test_img = np.clip(test_img.astype(int) + noise, 0, 255).astype(np.uint8)
    
    # 保存測試圖像
    Image.fromarray(test_img).save(test_image_path)
    print(f"✅ 創建測試圖像: {test_image_path}")
    
    # 測試1: 直接使用save_superpixel_regions函數
    print("\n📋 測試1: 直接測試save_superpixel_regions函數")
    
    # 轉換為浮點數並應用SLIC
    img_array = test_img.astype(np.float32) / 255.0
    segments = slic(img_array, n_segments=20, compactness=10, start_label=1, enforce_connectivity=True)
    
    # 保存超像素區域
    regions_dir = os.path.join(test_dir, "direct_test_regions")
    saved_regions = save_superpixel_regions(test_img, segments, regions_dir, "test_image", max_regions=20)
    
    print(f"✅ 直接測試完成，保存了 {len(saved_regions)} 個區域")
    
    # 測試2: 通過compute_decision_graph函數測試
    print("\n📋 測試2: 通過compute_decision_graph函數測試")
    
    result = compute_decision_graph(
        test_image_path,
        test_dir,
        use_superpixels=True,
        n_segments=15,
        compactness=10,
        save_regions=True,
        max_regions=15
    )
    
    print(f"✅ compute_decision_graph測試完成")
    
    # 測試3: 使用真實圖像（如果存在）
    print("\n📋 測試3: 使用真實圖像測試（如果存在）")
    
    # 嘗試找到一個真實圖像
    possible_image_paths = [
        "../Data/A01.jpg",
        "sample_image_1.png",
        "../sample_image_1.png"
    ]
    
    real_image_path = None
    for path in possible_image_paths:
        if os.path.exists(path):
            real_image_path = path
            break
    
    if real_image_path:
        print(f"🖼️  找到真實圖像: {real_image_path}")
        try:
            # 測試真實圖像
            real_result = compute_decision_graph(
                real_image_path,
                test_dir,
                use_superpixels=True,
                n_segments=50,
                compactness=10,
                save_regions=True,
                max_regions=30  # 限制保存數量
            )
            print(f"✅ 真實圖像測試完成")
        except Exception as e:
            print(f"❌ 真實圖像測試失敗: {e}")
    else:
        print("⚠️  未找到真實圖像，跳過此測試")
    
    # 顯示結果統計
    print(f"\n📊 測試結果統計:")
    print(f"測試目錄: {test_dir}")
    
    # 統計生成的文件
    total_files = 0
    for root, dirs, files in os.walk(test_dir):
        for file in files:
            if file.endswith('.png'):
                total_files += 1
    
    print(f"總共生成的PNG文件數量: {total_files}")
    
    # 列出所有子目錄
    subdirs = [d for d in os.listdir(test_dir) if os.path.isdir(os.path.join(test_dir, d))]
    if subdirs:
        print(f"生成的子目錄: {subdirs}")
        for subdir in subdirs:
            subdir_path = os.path.join(test_dir, subdir)
            files_in_subdir = [f for f in os.listdir(subdir_path) if f.endswith('.png')]
            print(f"  {subdir}: {len(files_in_subdir)} 個PNG文件")
    
    print(f"\n🎉 超像素區域保存功能測試完成！")
    print(f"📁 請查看 {test_dir} 目錄中的結果")

if __name__ == "__main__":
    test_superpixel_regions_save()