#!/usr/bin/env python3
"""
测试 Step 3 (precompute template matches) 的脚本
用法: python test_step3.py --dataset_name Task01_BrainTumour
"""

import sys
import os
from pathlib import Path

# 添加当前目录到路径
SCRIPT_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(SCRIPT_DIR))

from end_to_end_preprocessing import run_step3_precompute
import argparse

def main():
    parser = argparse.ArgumentParser(description="Test Step 3 (precompute template matches)")
    parser.add_argument("--dataset_name", type=str, required=True,
                       help="Dataset name (e.g., Task01_BrainTumour)")
    parser.add_argument("--output_base_dir", type=str, default="/projects/p32685/RL_dataset",
                       help="Base output directory")
    parser.add_argument("--template_percent", type=int, default=10,
                       help="Template percentage (default: 10)")
    parser.add_argument("--device", type=str, default="cuda:0",
                       help="Device for template precomputation (default: cuda:0)")
    parser.add_argument("--topk", type=int, default=1,
                       help="Top-k matches per sample (default: 1)")
    
    args = parser.parse_args()
    
    # 确定 processed 目录（Step 2 的输出）
    if args.dataset_name.startswith("Task") and args.dataset_name.endswith(("BrainTumour", "Heart", "Liver", "Hippocampus", 
                                                                          "Prostate", "Lung", "Pancreas", "HepaticVessel", 
                                                                          "Spleen", "Colon")):
        task_num = args.dataset_name.split("_")[0]  # Task01
        processed_dir = os.path.join(args.output_base_dir, f"MSD_{task_num.lower()}", 
                                   f"{args.dataset_name}_{args.template_percent}percent_processed")
    else:
        processed_dir = os.path.join(args.output_base_dir, 
                                   f"{args.dataset_name.lower()}_{args.template_percent}percent_processed")
    
    # 检查 processed 目录是否存在
    if not os.path.exists(processed_dir):
        print(f"✗ Error: Processed directory not found: {processed_dir}")
        print(f"  Make sure Step 2 has completed successfully.")
        sys.exit(1)
    
    # 检查必要的子目录
    hf_dataset_path = os.path.join(processed_dir, "hf_dataset")
    template_dir = os.path.join(processed_dir, "template")
    backup_dir = os.path.join(processed_dir, "backup_entries")
    
    missing_dirs = []
    if not os.path.isdir(hf_dataset_path):
        missing_dirs.append(f"hf_dataset ({hf_dataset_path})")
    if not os.path.isdir(template_dir):
        missing_dirs.append(f"template ({template_dir})")
    if not os.path.isdir(backup_dir):
        missing_dirs.append(f"backup_entries ({backup_dir})")
    
    if missing_dirs:
        print(f"✗ Error: Missing required directories:")
        for d in missing_dirs:
            print(f"  - {d}")
        print(f"  Make sure Step 2 has completed successfully.")
        sys.exit(1)
    
    print("="*80)
    print("TESTING STEP 3: PRECOMPUTE TEMPLATE MATCHES")
    print("="*80)
    print(f"Dataset: {args.dataset_name}")
    print(f"Processed directory: {processed_dir}")
    print(f"Device: {args.device}")
    print(f"Top-k matches: {args.topk}")
    print("="*80)
    
    try:
        run_step3_precompute(
            processed_dir=processed_dir,
            device=args.device,
            topk=args.topk
        )
        print("\n" + "="*80)
        print("✓ STEP 3 TEST PASSED")
        print("="*80)
        print(f"Processed directory: {processed_dir}")
        return 0
    except Exception as e:
        print("\n" + "="*80)
        print("✗ STEP 3 TEST FAILED")
        print("="*80)
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()
        return 1

if __name__ == "__main__":
    sys.exit(main())