import sys
import yaml
from pathlib import Path
from typing import Set, Tuple
import argparse  # <-- [新增]
import logging # <-- [新增]

try:
    # Import config to get base directories and dataset configuration path
    import yolo_config as config
except ImportError:
    print("错误：无法导入 'yolo_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 get_source_labels_dir() -> Path:
    """
    从 dataset.yaml 解析并构建源标签目录的绝对路径。
    路径的最后一部分 (例如 'train', 'val' 或 'test') 将根据 dataset.yaml 中的 'test' 键动态确定。
    """
    try:
        with open(config.DATASET_YAML_PATH, 'r', encoding='utf-8') as f:
            yaml_data = yaml.safe_load(f)
        
        # 1. 获取数据集的根目录相对路径
        relative_path_from_yaml = yaml_data.get('path')
        if not relative_path_from_yaml:
            print(f"错误: 在 '{config.DATASET_YAML_PATH}' 文件中没有找到 'path' 键。")
            sys.exit(1)

        # 2. 【新增】获取 'test' 键的值 (例如 'images/val')
        test_path_from_yaml = yaml_data.get('test')
        if not test_path_from_yaml:
            print(f"错误: 在 '{config.DATASET_YAML_PATH}' 文件中没有找到 'test' 键。")
            sys.exit(1)
        
        # 3. 【新增】从 'images/val' 中提取最后一部分 'val'
        #    Path(...).name 可以轻松实现这个功能
        source_split = Path(test_path_from_yaml).name

        # 4. 构建数据集的绝对路径
        dataset_dir = (config.DATASET_YAML_PATH.parent / relative_path_from_yaml).resolve()
        
        # 5. 【已修改】使用动态提取的 source_split 构建并返回标签目录路径
        labels_dir = dataset_dir / "labels" / source_split
        return labels_dir

    except FileNotFoundError:
        print(f"错误: YAML 配置文件未找到: '{config.DATASET_YAML_PATH}'")
        sys.exit(1)
    except Exception as e:
        print(f"读取或解析 YAML 文件时发生错误: {e}")
        sys.exit(1)


def main():
    """
    主函数，用于详细检查源训练标签集和所有预测输出的原始掩码
    目录之间的文件数量和文件名一致性。
    """
    # --- [新增] ---
    parser = argparse.ArgumentParser(description="检查预测掩码与源标签的一致性。")
    parser.add_argument(
        "--pt_name",
        type=str,
        default="best.pt",
        help="要检查的权重文件名称 (例如 'best.pt' 或 'epoch100.pt')。脚本将查找以此为后缀的掩码文件。"
    )
    args = parser.parse_args()

    # 从 'best.pt' 提取 'best'
    pt_name_raw = Path(args.pt_name).stem
    # 构建预期的后缀, e.g., '_best'
    expected_suffix = f"_{pt_name_raw}"

    # [核心修改] 只有 'best.pt' 会兼容检查无后缀的传统文件
    check_traditional = (args.pt_name == "best.pt")
    
    print(f"--- 开始详细检查预测原始掩码 (predicted_raw_masks) ---")
    if check_traditional:
        print(f"[i] 本次检查目标: --pt_name='{args.pt_name}' (检查后缀 '{expected_suffix}' 或 无后缀的传统文件)")
    else:
        print(f"[i] 本次检查目标: --pt_name='{args.pt_name}' (仅检查后缀 '{expected_suffix}' 的文件)")
    
    # 1. 获取源标签目录并解析文件信息
    source_dir = get_source_labels_dir()
    if not source_dir.exists() or not source_dir.is_dir():
        print(f"错误：在 '{source_dir}' 找不到源训练标签目录。请检查 dataset.yaml 中的 'path' 设置。")
        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)

    # 在你的日志中，路径是 '1_CholecSeg8k-13Type-1920x1080'，但配置文件是 '4_Dresden-11Type-512x512'
    # 这里我们以代码逻辑为准，它会动态读取配置文件
    print(f"源标签集: 在 '{source_dir}' 中找到 {source_file_count} 个标签文件 (.txt)。")
    print("-" * 60)

    # 2. 获取预测结果的基础目录
    predictions_base_dir = 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
    
    # 【已修正】直接查找所有模型的运行目录，例如 'YOLOv8n-seg_2025-09-20_10-00-00'
    run_dirs = sorted([d for d in predictions_base_dir.iterdir() if d.is_dir()])
    
    if not run_dirs:
        print(f"信息：在基础目录 '{predictions_base_dir}' 中没有找到任何模型的预测结果文件夹。")
        sys.exit(0)

    # 【已修正】移除外层循环，直接遍历 run_dirs
    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
        
        # 获取目录中所有的stems
        all_predicted_stems, _ = get_file_info(target_dir)
        
        # --- [核心修改] ---
        # 1. 提取我们关心的stems
        processed_stems = set()     # 存储处理后 (去掉后缀) 的stems, e.g., {'fileA', 'fileB'}
        stems_we_checked = set()    # 存储我们检查过的原始stems, e.g., {'fileA_best', 'fileA'}

        for stem in all_predicted_stems:
            # 检查 1: 是否为带后缀的新模式 (e.g., "fileA_best" 或 "fileA_epoch100")
            if stem.endswith(expected_suffix):
                # 去掉后缀, e.g., 'fileA_best' -> 'fileA'
                unstripped_stem = stem[:-len(expected_suffix)]
                
                # 仅当去掉后缀后的部分在源文件中时才添加
                if unstripped_stem in source_stems:
                    processed_stems.add(unstripped_stem)
                    stems_we_checked.add(stem)
            
            # 检查 2: [条件] 是否允许检查传统文件 (check_traditional == True)
            #         [条件] 且该文件是否为无后缀的传统文件 (e.g., "fileA")
            elif check_traditional and (stem in source_stems):
                # 兼容 'best.pt' 检查传统文件
                processed_stems.add(stem)
                stems_we_checked.add(stem)
        
        # 2. 现在，'processed_stems' 只包含与 'source_stems' 对应的基础文件名
        predicted_file_count = len(processed_stems)

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

        # --- 检查 2: 文件名 (基于处理后的stems) ---
        missing_files = source_stems - processed_stems
        if missing_files:
            examples = list(missing_files)[:3]
            print(f"  [✗ 文件名缺失] 预测结果中缺少 {len(missing_files)} 个文件 (匹配检查规则)。例如: {examples}...")
            is_current_dir_ok = False

        extra_files = processed_stems - source_stems
        if extra_files:
            examples = list(extra_files)[:3]
            print(f"  [✗ 文件名多余] (逻辑异常) 预测结果中多出 {len(extra_files)} 个文件。例如: {examples}...")
            is_current_dir_ok = False
        
        # --- [新增] 检查 3: 提示其他被忽略的文件 ---
        other_files = all_predicted_stems - stems_we_checked
        if other_files:
            examples = list(other_files)[:3]
            print(f"  [i 信息] 目录中还包含 {len(other_files)} 个其他文件 (例如: {examples}...)。")
            if check_traditional:
                print(f"  [i 信息] (这些文件不匹配后缀 '{expected_suffix}' 且不属于无后缀传统文件，已忽略)")
            else:
                 print(f"  [i 信息] (这些文件不匹配后缀 '{expected_suffix}'，已忽略)")

        # --- 当前目录的检查总结 ---
        if is_current_dir_ok:
            # 再次检查目录为空的特殊情况
            if predicted_file_count == 0 and source_file_count > 0:
                 print(f"  [✗ 目录为空] 预测目录中没有找到任何匹配本次检查规则的文件。")
                 global_mismatch_found = True
            else:
                 print(f"  [✓ OK] 所有检查通过 (文件数量: {predicted_file_count}，匹配检查规则)")
        else:
            global_mismatch_found = True

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

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


if __name__ == "__main__":
    main()