#!/usr/bin/env python3
"""
测试 Step 2 (packing) 的脚本
用法: python test_step2.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_step2_packing
import argparse

def main():
    parser = argparse.ArgumentParser(description="Test Step 2 (packing)")
    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)")
    
    args = parser.parse_args()
    
    # 确定输入目录（HF 格式）
    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
        dataset_dir = os.path.join(args.output_base_dir, f"MSD_{task_num.lower()}", args.dataset_name)
    else:
        dataset_dir = os.path.join(args.output_base_dir, args.dataset_name.lower())
    
    # 检查输入目录是否存在
    if not os.path.exists(dataset_dir):
        print(f"✗ Error: Dataset directory not found: {dataset_dir}")
        print(f"  Make sure Step 1 has completed successfully.")
        sys.exit(1)
    
    # 检查 dataset_info.json 是否存在
    dataset_info_path = os.path.join(dataset_dir, "dataset_info.json")
    if not os.path.exists(dataset_info_path):
        print(f"✗ Error: dataset_info.json not found in {dataset_dir}")
        print(f"  Step 1 may not have completed successfully.")
        sys.exit(1)
    
    print("="*80)
    print("TESTING STEP 2: PACKING")
    print("="*80)
    print(f"Dataset: {args.dataset_name}")
    print(f"Input directory: {dataset_dir}")
    print(f"Output base: {args.output_base_dir}")
    print(f"Template percent: {args.template_percent}%")
    print("="*80)
    
    try:
        processed_dir = run_step2_packing(
            dataset_dir=dataset_dir,
            output_base_dir=args.output_base_dir,
            dataset_name=args.dataset_name,
            template_percent=args.template_percent
        )
        print("\n" + "="*80)
        print("✓ STEP 2 TEST PASSED")
        print("="*80)
        print(f"Output directory: {processed_dir}")
        return 0
    except Exception as e:
        print("\n" + "="*80)
        print("✗ STEP 2 TEST FAILED")
        print("="*80)
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()
        return 1

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