#!/usr/bin/env python3
"""
Batch process four BiomedParse datasets to RL_dataset format.

Datasets:
- BraTS2020
- SPIDER
- STS-2D-Tooth
- STS-3D-Tooth

Input:  /projects/p32685/datasets/biomedparse/{BraTS2020,SPIDER,STS-2D-Tooth,STS-3D-Tooth}
Output: /projects/p32685/RL_dataset/{brats2020,spider,sts-2d-tooth,sts-3d-tooth}_10percent_processed/
"""

import os
import sys
import subprocess
from pathlib import Path

SCRIPT_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(SCRIPT_DIR))

FOUR_DATASETS = [
    {
        "name": "BraTS2020",
        "raw_dir": "/projects/p32685/datasets/biomedparse/BraTS2020",
        "description": "BraTS 2020 - Brain Tumour Segmentation"
    },
    {
        "name": "SPIDER",
        "raw_dir": "/projects/p32685/datasets/biomedparse/SPIDER",
        "description": "SPIDER - Spine Segmentation Dataset"
    },
    {
        "name": "STS-2D-Tooth",
        "raw_dir": "/projects/p32685/datasets/biomedparse/STS-2D-Tooth",
        "description": "STS-2D-Tooth - 2D Tooth Segmentation Dataset"
    },
    {
        "name": "STS-3D-Tooth",
        "raw_dir": "/projects/p32685/datasets/biomedparse/STS-3D-Tooth",
        "description": "STS-3D-Tooth - 3D Tooth Segmentation Dataset"
    }
]

OUTPUT_BASE_DIR = "/projects/p32685/RL_dataset"
TEMPLATE_PERCENT = 10
DEVICE = "cuda:0"


def process_dataset(dataset_config, skip_convert=False, skip_packing=False, skip_precompute=False):
    """Process a single dataset."""
    dataset_name = dataset_config["name"]
    raw_dir = dataset_config["raw_dir"]

    print("\n" + "="*80)
    print(f"Processing {dataset_name}: {dataset_config['description']}")
    print("="*80)
    print(f"Raw data: {raw_dir}")
    print(f"Output base: {OUTPUT_BASE_DIR}")
    print(f"Template percent: {TEMPLATE_PERCENT}%")

    if not os.path.exists(raw_dir):
        print(f"✗ Error: Raw data directory does not exist: {raw_dir}")
        return False

    cmd = [
        "python3",
        os.path.join(SCRIPT_DIR, "end_to_end_preprocessing.py"),
        "--dataset_name", dataset_name,
        "--template_percent", str(TEMPLATE_PERCENT),
        "--raw_data_dir", raw_dir,
        "--output_base_dir", OUTPUT_BASE_DIR,
        "--device", DEVICE
    ]
    if skip_convert:
        cmd.append("--skip_convert")
    if skip_packing:
        cmd.append("--skip_packing")
    if skip_precompute:
        cmd.append("--skip_precompute")

    print(f"\nRunning: {' '.join(cmd)}")

    try:
        result = subprocess.run(cmd, check=True, cwd=SCRIPT_DIR)
        if result.returncode == 0:
            print(f"\n✓ {dataset_name} processed successfully!")
            return True
        else:
            print(f"\n✗ {dataset_name} failed with return code {result.returncode}")
            return False
    except subprocess.CalledProcessError as e:
        print(f"\n✗ {dataset_name} failed: {e}")
        return False
    except Exception as e:
        print(f"\n✗ {dataset_name} failed with exception: {e}")
        import traceback
        traceback.print_exc()
        return False


def main():
    import argparse

    parser = argparse.ArgumentParser(description="Batch process four BiomedParse datasets to RL_dataset")
    parser.add_argument("--start_from", type=int, default=1, help="Start from dataset number (1-4)")
    parser.add_argument("--end_at", type=int, default=4, help="End at dataset number (1-4)")
    parser.add_argument("--skip_convert", action="store_true", help="Skip step 1 (conversion)")
    parser.add_argument("--skip_packing", action="store_true", help="Skip step 2 (packing)")
    parser.add_argument("--skip_precompute", action="store_true", help="Skip step 3 (precompute)")
    parser.add_argument("--device", type=str, default="cuda:0", help="Device for template precomputation")
    parser.add_argument("--template_percent", type=int, default=10, help="Template percentage (default: 10)")

    args = parser.parse_args()
    global DEVICE, TEMPLATE_PERCENT
    DEVICE = args.device
    TEMPLATE_PERCENT = args.template_percent

    print("╔" + "="*78 + "╗")
    print("║" + " "*15 + "BATCH PROCESS BIOMEDPARSE FOUR DATASETS" + " "*26 + "║")
    print("╔" + "="*78 + "╗")
    print(f"\nDatasets: BraTS2020, SPIDER, STS-2D-Tooth, STS-3D-Tooth")
    print(f"Processing: {args.start_from} to {args.end_at}")
    print(f"Output base: {OUTPUT_BASE_DIR}")
    print(f"Template percent: {TEMPLATE_PERCENT}%")
    print(f"Device: {DEVICE}")

    if args.skip_convert:
        print("⚠ Skipping Step 1 (conversion)")
    if args.skip_packing:
        print("⚠ Skipping Step 2 (packing)")
    if args.skip_precompute:
        print("⚠ Skipping Step 3 (precompute)")

    success_count = 0
    failed = []

    for i, cfg in enumerate(FOUR_DATASETS, start=1):
        if i < args.start_from or i > args.end_at:
            continue
        print(f"\n{'='*80}")
        print(f"Dataset {i}/4: {cfg['name']}")
        print(f"{'='*80}")
        if process_dataset(cfg, args.skip_convert, args.skip_packing, args.skip_precompute):
            success_count += 1
        else:
            failed.append(cfg['name'])

    print("\n" + "="*80)
    print("BATCH PROCESSING SUMMARY")
    print("="*80)
    n_total = args.end_at - args.start_from + 1
    print(f"Successfully processed: {success_count}/{n_total} datasets")
    if failed:
        print(f"Failed: {', '.join(failed)}")
    else:
        print("\n✓ All datasets processed successfully!")
    print(f"\nOutput: {OUTPUT_BASE_DIR}/")
    for i in range(args.start_from, args.end_at + 1):
        name = FOUR_DATASETS[i-1]['name']
        out = f"{name.lower()}_{TEMPLATE_PERCENT}percent_processed"
        print(f"  {out}/")


if __name__ == "__main__":
    main()