# predict_check.py

import sys
from pathlib import Path
from typing import Set, Tuple

try:
    import config
except ImportError:
    print("错误：无法导入 'config.py'。请确保此脚本与 config.py 在同一目录下。")
    sys.exit(1)

def get_file_info(directory: Path) -> Tuple[Set[str], Set[str]]:
    """
    安全地获取目录中所有文件的基本名（stem）和后缀（suffix）。
    
    返回:
        一个包含 stems 集合和 suffixes 集合的元组。
    """
    if not directory.is_dir():
        return set(), set()
    
    stems = set()
    suffixes = set()
    for item in directory.iterdir():
        if item.is_file():
            stems.add(item.stem)
            suffixes.add(item.suffix)
    return stems, suffixes

def main():
    """
    主函数，用于详细检查源测试数据集和所有预测输出目录之间的
    文件数量、文件名和后缀一致性。
    """
    print("--- 开始详细检查预测输出 ---")

    # 1. 从 config.py 获取源目录并解析文件信息
    source_dir = config.TEST_IMAGE_DIR
    if not source_dir.exists() or not source_dir.is_dir():
        print(f"错误：在 '{source_dir}' 找不到源测试图片目录。")
        sys.exit(1)

    source_stems, _ = get_file_info(source_dir)
    source_file_count = len(source_stems)
    
    if source_file_count == 0:
        print(f"警告：源目录 '{source_dir}' 为空或不包含任何文件。检查中止。")
        sys.exit(0)

    print(f"源数据集: 在 '{source_dir}' 中找到 {source_file_count} 张图片。")
    print("-" * 50)

    # 2. 从 config.py 获取预测结果的基础目录
    predictions_base_dir = Path(config.PREDICT_BEST_MODEL_DIR)
    if not predictions_base_dir.exists() or not predictions_base_dir.is_dir():
        print(f"错误：在 '{predictions_base_dir}' 找不到预测结果的基础目录。")
        sys.exit(1)

    # 3. 遍历每个模型的运行目录并进行详细检查
    found_predictions = False
    global_mismatch_found = False
    
    run_dirs = sorted([d for d in predictions_base_dir.iterdir() if d.is_dir()])

    if not run_dirs:
        print("信息：在基础目录中没有找到任何模型的运行记录文件夹。")
        sys.exit(0)

    for run_dir in run_dirs:
        target_dir = run_dir / "predicted_raw_masks"
        
        # 如果目标目录不存在，则跳过
        if not (target_dir.exists() and target_dir.is_dir()):
            continue

        print(f"正在检查: '{target_dir}'...")
        found_predictions = True
        is_current_dir_ok = True
        
        predicted_stems, predicted_suffixes = get_file_info(target_dir)
        predicted_file_count = len(predicted_stems)

        # --- 检查 1: 文件数量 ---
        if source_file_count != predicted_file_count:
            print(f"  [✗ 数量不匹配] 源目录有 {source_file_count} 个文件，但此目录有 {predicted_file_count} 个。")
            is_current_dir_ok = False

        # --- 检查 2: 文件名 ---
        missing_files = source_stems - predicted_stems
        if missing_files:
            # 只显示最多3个示例文件名以保持简洁
            examples = list(missing_files)[:3]
            print(f"  [✗ 文件名缺失] 预测结果中缺少 {len(missing_files)} 个文件。例如: {examples}...")
            is_current_dir_ok = False

        extra_files = predicted_stems - source_stems
        if extra_files:
            examples = list(extra_files)[:3]
            print(f"  [✗ 文件名多余] 预测结果中多出 {len(extra_files)} 个文件。例如: {examples}...")
            is_current_dir_ok = False

        # --- 检查 3: 后缀一致性 ---
        # 只有在目录不为空的情况下，后缀多于1个才算问题
        if predicted_file_count > 0 and len(predicted_suffixes) > 1:
            print(f"  [✗ 后缀不一致] 预测目录中存在多种文件后缀: {sorted(list(predicted_suffixes))}")
            is_current_dir_ok = False
        
        # --- 当前目录的检查总结 ---
        if is_current_dir_ok:
            # 再次检查目录为空的特殊情况
            if predicted_file_count == 0 and source_file_count > 0:
                 print("  [✗ 目录为空] 预测目录是空的，但源目录包含文件。")
                 global_mismatch_found = True
            else:
                 suffix_info = list(predicted_suffixes)[0] if len(predicted_suffixes) == 1 else 'N/A'
                 print(f"  [✓ OK] 所有检查通过 (数量: {predicted_file_count}, 后缀: {suffix_info})")
        else:
            global_mismatch_found = True

        print("-" * 25)

    # 4. 输出最终的全局检查摘要
    print("\n--- 检查摘要 ---")
    if not found_predictions:
        print("结果: 没有找到任何 'predicted_raw_masks' 目录进行检查。")
    elif global_mismatch_found:
        print("结论: 检查不通过。发现至少一个预测目录未通过所有检查，请查看上方日志。")
        sys.exit(1)
    else:
        print("结论: 检查通过！所有已找到的预测目录均通过了数量、文件名和后缀一致性检查。")

    print("--- 检查完成 ---")


if __name__ == "__main__":
    main()