#!/usr/bin/env python3
"""
Process lungx BiomedParse dataset to RL_dataset format.

Dataset: lungx (Montgomery + Shenzhen Chest X-Ray lung segmentation)
  - Classes: background(0), left_lung(1), right_lung(2)
  - Train: 110 images, Test: 28 images
  - Masks: binary per-class ({base}_left_lung.png, {base}_right_lung.png)

Input:  /projects/p32685/datasets/biomedparse/lungx
Output: /projects/p32685/RL_dataset/lungx/
        /projects/p32685/RL_dataset/lungx_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))

DATASET_CONFIG = {
    "name": "lungx",
    "raw_dir": "/projects/p32685/datasets/biomedparse/lungx",
    "description": "Lung X-Ray - Left and Right Lung Segmentation (Montgomery + Shenzhen)",
}

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):
    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 = [
        sys.executable,
        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="Process lungx BiomedParse dataset to RL_dataset")
    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 + "PROCESS LUNGX BIOMEDPARSE DATASET" + " " * 30 + "║")
    print("╔" + "=" * 78 + "╗")
    print(f"\nDataset: lungx")
    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 = process_dataset(
        DATASET_CONFIG,
        args.skip_convert,
        args.skip_packing,
        args.skip_precompute,
    )

    print("\n" + "=" * 80)
    print("PROCESSING SUMMARY")
    print("=" * 80)
    if success:
        print("✓ lungx processed successfully!")
        out = f"lungx_{TEMPLATE_PERCENT}percent_processed"
        print(f"\nOutput: {OUTPUT_BASE_DIR}/")
        print(f"  lungx/                    (converted HF format)")
        print(f"  {out}/   (processed with templates)")
    else:
        print("✗ lungx processing failed!")


if __name__ == "__main__":
    main()