#!/usr/bin/env python3
"""Script to run quality check workflow in separate rounds for memory management.

This script allows running different types of image generation in separate rounds,
so that memory can be cleared between rounds. It also supports generating atlas views
once for all subjects instead of per subject.
"""

from __future__ import annotations

import argparse
import gc
import json
import logging
import os
import sys
from pathlib import Path

# Try to import fcntl (Unix/Linux only)
try:
    import fcntl
    HAS_FCNTL = True
except ImportError:
    HAS_FCNTL = False

# Add src to path so we can import pytractoviz
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

from pytractoviz.html import create_quality_check_html
from pytractoviz.viz import TractographyVisualizer

logger = logging.getLogger(__name__)


def load_config(config_file: Path) -> dict:
    """Load configuration from JSON file."""
    with open(config_file) as f:
        return json.load(f)


def save_results(results: dict, output_file: Path) -> None:
    """Save results dictionary to JSON file with file locking for thread safety.
    
    Uses atomic write (write to temp file then rename) to prevent corruption
    when multiple processes write simultaneously.
    """
    # Convert Path objects to strings for JSON serialization
    results_serializable: dict[str, dict[str, dict[str, str]]] = {}
    for subject_id, subject_tracts in results.items():
        results_serializable[subject_id] = {}
        if isinstance(subject_tracts, dict):
            for tract_name, media_dict in subject_tracts.items():
                results_serializable[subject_id][tract_name] = {}
                if isinstance(media_dict, dict):
                    for media_type, file_path in media_dict.items():
                        results_serializable[subject_id][tract_name][media_type] = str(file_path)

    # Atomic write: write to temp file, then rename
    output_file = Path(output_file)
    output_file.parent.mkdir(parents=True, exist_ok=True)

    # Create temp file in same directory for atomic rename
    temp_file = output_file.with_suffix(output_file.suffix + ".tmp")

    try:
        with open(temp_file, "w") as f:
            json.dump(results_serializable, f, indent=2)
        # Atomic rename (works on Unix/Linux, Windows may need different approach)
        temp_file.replace(output_file)
    except Exception:
        # Clean up temp file on error
        if temp_file.exists():
            temp_file.unlink()
        raise


def load_results(results_file: Path) -> dict[str, dict[str, dict[str, str]]]:
    """Load results dictionary from JSON file with file locking for thread safety."""
    results_file = Path(results_file)
    if not results_file.exists():
        return {}

    # Use file locking to prevent reading while another process is writing
    if HAS_FCNTL:
        try:
            with open(results_file) as f:
                # Try to acquire shared lock (non-blocking)
                try:
                    fcntl.flock(f.fileno(), fcntl.LOCK_SH | fcntl.LOCK_NB)
                except OSError:
                    # If we can't get lock, wait a bit and try again
                    # This handles the case where another process is writing
                    import time
                    time.sleep(0.1)
                    fcntl.flock(f.fileno(), fcntl.LOCK_SH)

                try:
                    return json.load(f)
                finally:
                    fcntl.flock(f.fileno(), fcntl.LOCK_UN)
        except OSError as e:
            # If file locking fails, fall back to regular read
            logger.warning("File locking failed, using regular read: %s", e)
            with open(results_file) as f:
                return json.load(f)
    else:
        # No fcntl available (Windows), use regular read
        # Note: This is not thread-safe but better than failing
        with open(results_file) as f:
            return json.load(f)


def merge_results(existing: dict, new: dict) -> dict:
    """Merge new results into existing results."""
    for subject_id, subject_tracts in new.items():
        if subject_id not in existing:
            existing[subject_id] = {}
        if isinstance(subject_tracts, dict):
            for tract_name, media_dict in subject_tracts.items():
                if tract_name not in existing[subject_id]:
                    existing[subject_id][tract_name] = {}
                if isinstance(media_dict, dict):
                    existing[subject_id][tract_name].update(media_dict)
    return existing


def save_and_merge_results(
    new_results: dict,
    results_file: Path,
    max_retries: int = 10,
    retry_delay: float = 0.1,
) -> None:
    """Atomically merge new results into existing results file.
    
    This function handles concurrent writes by:
    1. Loading existing results with file lock
    2. Merging new results
    3. Saving with atomic write
    
    Parameters
    ----------
    new_results : dict
        New results to merge.
    results_file : Path
        Path to results JSON file.
    max_retries : int, optional
        Maximum number of retries if file is locked. Default is 10.
    retry_delay : float, optional
        Delay between retries in seconds. Default is 0.1.
    """
    results_file = Path(results_file)

    for attempt in range(max_retries):
        try:
            # Load existing results with lock
            existing_results = load_results(results_file)

            # Merge results
            merged_results = merge_results(existing_results.copy(), new_results)

            # Save with atomic write
            save_results(merged_results, results_file)

            # Success!
            return
        except OSError as e:
            if attempt < max_retries - 1:
                # Wait and retry
                import time
                time.sleep(retry_delay * (attempt + 1))  # Exponential backoff
                logger.debug(
                    "Retry %d/%d: File locked, waiting...",
                    attempt + 1,
                    max_retries,
                )
            else:
                # Last attempt failed
                logger.error("Failed to save results after %d attempts: %s", max_retries, e)
                raise


def split_subjects_into_batches(
    subjects: dict[str, dict[str, str | Path]],
    num_batches: int,
) -> list[dict[str, dict[str, str | Path]]]:
    """Split subjects dictionary into batches.
    
    Parameters
    ----------
    subjects : dict[str, dict[str, str | Path]]
        Dictionary mapping subject IDs to their tract files.
    num_batches : int
        Number of batches to create.
    
    Returns
    -------
    list[dict[str, dict[str, str | Path]]]
        List of subject dictionaries, one per batch.
    """
    subject_ids = sorted(subjects.keys())
    total_subjects = len(subject_ids)

    if num_batches <= 0:
        raise ValueError("num_batches must be positive")

    if num_batches >= total_subjects:
        # More batches than subjects - each subject gets its own batch
        return [{sid: subjects[sid]} for sid in subject_ids]

    # Calculate batch size
    batch_size = max(1, total_subjects // num_batches)
    batches: list[dict[str, dict[str, str | Path]]] = []

    for i in range(0, total_subjects, batch_size):
        batch_subject_ids = subject_ids[i : i + batch_size]
        batch = {sid: subjects[sid] for sid in batch_subject_ids}
        batches.append(batch)

    return batches


def get_slurm_array_task_id() -> int | None:
    """Get SLURM array task ID from environment.
    
    Returns
    -------
    int | None
        Array task ID (1-indexed) if running in SLURM array job, None otherwise.
    """
    array_task_id = os.environ.get("SLURM_ARRAY_TASK_ID")
    if array_task_id is not None:
        try:
            return int(array_task_id)
        except ValueError:
            logger.warning("Invalid SLURM_ARRAY_TASK_ID: %s", array_task_id)
            return None
    return None


def filter_subjects_by_batch(
    subjects: dict[str, dict[str, str | Path]],
    batch_index: int,
    num_batches: int,
) -> dict[str, dict[str, str | Path]]:
    """Filter subjects to only include those in the specified batch.
    
    Parameters
    ----------
    subjects : dict[str, dict[str, str | Path]]
        Dictionary mapping subject IDs to their tract files.
    batch_index : int
        Batch index (1-indexed, matching SLURM_ARRAY_TASK_ID).
    num_batches : int
        Total number of batches.
    
    Returns
    -------
    dict[str, dict[str, str | Path]]
        Filtered subjects dictionary containing only subjects in the specified batch.
    """
    batches = split_subjects_into_batches(subjects, num_batches)

    # batch_index is 1-indexed (from SLURM), convert to 0-indexed
    if batch_index < 1 or batch_index > len(batches):
        raise ValueError(
            f"Batch index {batch_index} is out of range (1-{len(batches)})",
        )

    return batches[batch_index - 1]


def filter_subjects_by_list(
    subjects: dict[str, dict[str, str | Path]],
    subject_list: list[str],
) -> dict[str, dict[str, str | Path]]:
    """Filter subjects to only include those in the provided list.
    
    Parameters
    ----------
    subjects : dict[str, dict[str, str | Path]]
        Dictionary mapping subject IDs to their tract files.
    subject_list : list[str]
        List of subject IDs to include.
    
    Returns
    -------
    dict[str, dict[str, str | Path]]
        Filtered subjects dictionary containing only subjects in the list.
    """
    filtered = {
        subject_id: subject_tracts
        for subject_id, subject_tracts in subjects.items()
        if subject_id in subject_list
    }

    # Warn about missing subjects
    missing = set(subject_list) - set(subjects.keys())
    if missing:
        logger.warning("Subjects not found in config: %s", ", ".join(sorted(missing)))

    return filtered


def load_subject_list(subject_file: Path) -> list[str]:
    """Load subject IDs from a file (one per line).
    
    Parameters
    ----------
    subject_file : Path
        Path to file containing subject IDs, one per line.
    
    Returns
    -------
    list[str]
        List of subject IDs.
    """
    with open(subject_file) as f:
        subjects = [line.strip() for line in f if line.strip() and not line.strip().startswith("#")]
    return subjects


def run_round(
    round_name: str,
    config: dict,
    results_file: Path,
    output_dir: Path,
    batch_index: int | None = None,
    num_batches: int | None = None,
) -> None:
    """Run a specific round of image generation.
    
    Parameters
    ----------
    round_name : str
        Name of the round to run.
    config : dict
        Configuration dictionary.
    results_file : Path
        Path to results JSON file.
    output_dir : Path
        Output directory.
    batch_index : int | None, optional
        Batch index (1-indexed) if running in batch mode. If None, processes all subjects.
    num_batches : int | None, optional
        Total number of batches. Required if batch_index is provided.
    """
    logger.info("=" * 80)
    logger.info("Starting round: %s", round_name)
    if batch_index is not None:
        logger.info("Processing batch %d of %d", batch_index, num_batches or 1)
    logger.info("=" * 80)

    # Load existing results
    existing_results = load_results(results_file)

    # Create visualizer
    visualizer = TractographyVisualizer(
        output_directory=str(output_dir),
        n_jobs=config.get("n_jobs", 1),
        figure_size=tuple(config.get("figure_size", [800, 800])),
    )

    # Determine which checks to skip
    all_checks = [
        "anatomical_views",
        "atlas_comparison",
        "cci",
        "before_after_cci",
        "afq_profile",
        "bundle_assignment",
        "shape_similarity",
    ]

    # Determine skip_checks based on round
    if round_name == "anatomical":
        skip_checks = [c for c in all_checks if c != "anatomical_views"]
    elif round_name == "cci":
        skip_checks = [c for c in all_checks if c != "cci"]
    elif round_name == "before_after_cci":
        skip_checks = [c for c in all_checks if c != "before_after_cci"]
    elif round_name == "atlas_views":
        # For atlas views, we'll generate them once for all subjects
        # We also generate subject MNI views here for atlas comparison
        skip_checks = all_checks  # Skip all other checks (handled separately)
    elif round_name == "shape_similarity":
        skip_checks = [c for c in all_checks if c != "shape_similarity"]
    elif round_name == "afq_profile":
        skip_checks = [c for c in all_checks if c != "afq_profile"]
    elif round_name == "bundle_assignment":
        skip_checks = [c for c in all_checks if c != "bundle_assignment"]
    else:
        raise ValueError(f"Unknown round name: {round_name}")

    # Prepare data structures
    subjects_original_space = config["subjects_original_space"]
    subjects_mni_space = config.get("subjects_mni_space")
    atlas_files = config.get("atlas_files")
    metric_files = config.get("metric_files")
    atlas_ref_img = config.get("atlas_ref_img")
    ref_img = config.get("ref_img")
    flip_lr = config.get("flip_lr", False)

    # Filter subjects by provided list first (if any)
    # If subjects are explicitly provided, skip batch filtering
    subject_list = config.get("subject_list")
    has_explicit_subjects = subject_list is not None

    if subject_list:
        logger.info("Filtering subjects by provided list (%d subjects)", len(subject_list))
        subjects_original_space = filter_subjects_by_list(
            subjects_original_space,
            subject_list,
        )
        logger.info("Processing %d subjects from list", len(subjects_original_space))

    # Filter subjects by batch if running in batch mode (only if no explicit subject list)
    if batch_index is not None and num_batches is not None and not has_explicit_subjects:
        logger.info("Filtering subjects for batch %d of %d", batch_index, num_batches)
        subjects_original_space = filter_subjects_by_batch(
            subjects_original_space,
            batch_index,
            num_batches,
        )
        logger.info("Processing %d subjects in this batch", len(subjects_original_space))
    elif batch_index is not None and has_explicit_subjects:
        logger.info("Explicit subject list provided, skipping batch filtering")

    # Also filter subjects_mni_space and metric_files if they exist
    if subjects_mni_space is not None:
        subjects_mni_space = {
            sid: tracts
            for sid, tracts in subjects_mni_space.items()
            if sid in subjects_original_space
        }

    if metric_files is not None:
        metric_files = {
            sid: metrics
            for sid, metrics in metric_files.items()
            if sid in subjects_original_space
        }

    # Filter ref_img if it's a dictionary
    if isinstance(ref_img, dict):
        ref_img = {
            sid: path
            for sid, path in ref_img.items()
            if sid in subjects_original_space
        }

    # Handle reference image - can be single path or dict mapping subject_id -> path
    # Convert string paths to Path objects, keep dict as-is (workflow handles it)
    if isinstance(ref_img, str):
        # Single reference image for all subjects
        ref_img_path: str | Path | dict[str, str | Path] | None = Path(ref_img)
    elif isinstance(ref_img, dict):
        # Per-subject reference images - convert values to Path objects
        ref_img_path = {subject_id: Path(path) for subject_id, path in ref_img.items()}
    else:
        ref_img_path = None

    # Special handling for atlas_views round - generate once for all subjects
    if round_name == "atlas_views":
        if atlas_files is None:
            logger.warning("No atlas_files provided, skipping atlas_views round")
            return

        # Create a shared atlas output directory
        atlas_output_dir = output_dir / "atlas_views"
        atlas_output_dir.mkdir(parents=True, exist_ok=True)

        # Generate atlas views once for each tract
        atlas_results: dict[str, dict[str, dict[str, str]]] = {}
        for tract_name, atlas_file in atlas_files.items():
            logger.info("Generating atlas views for tract: %s", tract_name)
            try:
                # Note: Set flip_lr=False to match MNI views (which should match original views)
                atlas_views = visualizer.generate_atlas_views(
                    atlas_file,
                    atlas_ref_img=atlas_ref_img,
                    flip_lr=False,  # Set to False to match MNI views
                    output_dir=atlas_output_dir,
                    **config.get("atlas_kwargs", {}),
                )

                # Add atlas views to results for all subjects that have this tract
                for subject_id, tracts in subjects_original_space.items():
                    if tract_name in tracts:
                        if subject_id not in atlas_results:
                            atlas_results[subject_id] = {}
                        if tract_name not in atlas_results[subject_id]:
                            atlas_results[subject_id][tract_name] = {}
                        # Add atlas views with the same keys as in the workflow
                        for view_name, view_path in atlas_views.items():
                            atlas_results[subject_id][tract_name][f"atlas_{view_name}"] = str(view_path)

                # Clean up after each atlas
                del atlas_views
                gc.collect()
            except Exception:
                logger.exception("Failed to generate atlas views for %s", tract_name)

        # Also generate subject MNI views if subjects_mni_space is provided
        # These are needed for atlas comparison even though atlas views are generated once
        if subjects_mni_space is not None:
            logger.info("Generating subject MNI views for atlas comparison")
            for subject_id, tracts in subjects_mni_space.items():
                subject_output_dir = output_dir / subject_id
                for tract_name, tract_file_mni in tracts.items():
                    if tract_name in atlas_files:  # Only generate if atlas exists
                        tract_output_dir = subject_output_dir / tract_name
                        tract_output_dir.mkdir(parents=True, exist_ok=True)
                        try:
                            # Note: Don't apply flip_lr here - MNI views should match original anatomical views
                            subject_mni_views = visualizer.generate_anatomical_views(
                                tract_file_mni,
                                ref_img=atlas_ref_img,  # Use atlas ref image for MNI space
                                output_dir=tract_output_dir,
                                **config.get("subject_kwargs", {}),
                            )
                            # Add subject MNI views to results
                            if subject_id not in atlas_results:
                                atlas_results[subject_id] = {}
                            if tract_name not in atlas_results[subject_id]:
                                atlas_results[subject_id][tract_name] = {}
                            for view_name, view_path in subject_mni_views.items():
                                atlas_results[subject_id][tract_name][f"subject_mni_{view_name}"] = str(view_path)
                            del subject_mni_views
                            gc.collect()
                        except Exception:
                            logger.exception(
                                "Failed to generate subject MNI views for %s/%s",
                                subject_id,
                                tract_name,
                            )

        # Merge atlas results into existing results atomically
        save_and_merge_results(atlas_results, results_file)
        logger.info("Atlas views round completed")
        return

    # For other rounds, run the normal workflow with skip_checks
    try:
        round_results = visualizer.run_quality_check_workflow(
            subjects_original_space=subjects_original_space,
            ref_img=ref_img_path,
            subjects_mni_space=subjects_mni_space,
            atlas_files=atlas_files,
            metric_files=metric_files,
            atlas_ref_img=atlas_ref_img,
            flip_lr=flip_lr,
            output_dir=output_dir,
            skip_checks=skip_checks,
            html_output=None,  # Don't generate HTML in individual rounds
            subject_kwargs=config.get("subject_kwargs"),
            atlas_kwargs=config.get("atlas_kwargs"),
            **config.get("kwargs", {}),
        )

        # Merge results atomically (handles concurrent writes)
        save_and_merge_results(round_results, results_file)

        logger.info("Round %s completed successfully", round_name)
    except Exception:
        logger.exception("Error in round %s", round_name)
        raise
    finally:
        # Clean up visualizer
        del visualizer
        gc.collect()


def generate_slurm_batch_script(
    config_file: Path,
    output_dir: Path,
    num_batches: int,
    rounds: list[str] | None = None,
    slurm_options: dict[str, str] | None = None,
    script_path: Path | None = None,
) -> Path:
    """Generate a SLURM batch script for array jobs.
    
    Parameters
    ----------
    config_file : Path
        Path to configuration JSON file.
    output_dir : Path
        Output directory.
    num_batches : int
        Number of batches (array job size).
    rounds : list[str] | None, optional
        List of rounds to run. If None, runs all rounds.
    slurm_options : dict[str, str] | None, optional
        Additional SLURM options (e.g., {"time": "24:00:00", "mem": "32G"}).
    script_path : Path | None, optional
        Path to save the batch script. If None, saves to output_dir/batch_submit.sh.
    
    Returns
    -------
    Path
        Path to the generated batch script.
    """
    if rounds is None:
        rounds = [
            "anatomical",
            "cci",
            "before_after_cci",
            "atlas_views",
            "shape_similarity",
            "afq_profile",
            "bundle_assignment",
        ]

    if slurm_options is None:
        slurm_options = {}

    if script_path is None:
        script_path = output_dir / "batch_submit.sh"

    # Get absolute paths
    config_file = config_file.resolve()
    output_dir = output_dir.resolve()
    run_rounds_script = Path(__file__).resolve()

    # Build SLURM directives
    slurm_directives = [
        "#!/bin/bash",
        f"#SBATCH --array=1-{num_batches}",
        "#SBATCH --job-name=pytractoviz_qc",
        f"#SBATCH --output={output_dir}/slurm-%A_%a.out",
        f"#SBATCH --error={output_dir}/slurm-%A_%a.err",
    ]

    # Add custom SLURM options
    for key, value in slurm_options.items():
        if key.startswith("#SBATCH"):
            slurm_directives.append(f"{key} {value}")
        else:
            slurm_directives.append(f"#SBATCH --{key}={value}")

    # Build the script content
    script_lines = slurm_directives + [
        "",
        "# Get array task ID (1-indexed)",
        "BATCH_ID=${SLURM_ARRAY_TASK_ID}",
        "",
        "# Run each round for this batch",
    ]

    for round_name in rounds:
        script_lines.extend([
            f"echo 'Running round: {round_name} (batch $BATCH_ID)'",
            f"python3 {run_rounds_script} {round_name} \\",
            f"    --config {config_file} \\",
            f"    --output-dir {output_dir} \\",
            "    --batch-index $BATCH_ID \\",
            f"    --num-batches {num_batches}",
            "",
            "# Check if round failed",
            "if [ $? -ne 0 ]; then",
            f"    echo 'ERROR: Round {round_name} failed for batch $BATCH_ID'",
            "    exit 1",
            "fi",
            "",
        ])

    script_lines.append("echo 'All rounds completed for batch $BATCH_ID'")

    # Write script
    script_path.parent.mkdir(parents=True, exist_ok=True)
    with open(script_path, "w") as f:
        f.write("\n".join(script_lines))

    # Make executable
    script_path.chmod(0o755)

    logger.info("Generated SLURM batch script: %s", script_path)
    logger.info("Submit with: sbatch %s", script_path)

    return script_path


def generate_html(config: dict, results_file: Path, output_dir: Path) -> None:
    """Generate HTML report from accumulated results."""
    logger.info("=" * 80)
    logger.info("Generating HTML report")
    logger.info("=" * 80)

    # Load all results
    results = load_results(results_file)

    if not results:
        logger.warning("No results found, cannot generate HTML report")
        return

    # Convert string paths back to Path objects for HTML generation
    results_for_html: dict[str, dict[str, dict[str, str]]] = {}
    for subject_id, subject_tracts in results.items():
        results_for_html[subject_id] = {}
        if isinstance(subject_tracts, dict):
            for tract_name, media_dict in subject_tracts.items():
                results_for_html[subject_id][tract_name] = {}
                if isinstance(media_dict, dict):
                    for media_type, file_path in media_dict.items():
                        results_for_html[subject_id][tract_name][media_type] = str(file_path)

    # Generate HTML
    html_output_config = config.get("html_output")
    html_output = output_dir / "quality_check_report.html" if html_output_config is None else Path(html_output_config)

    create_quality_check_html(
        results_for_html,
        str(html_output),
        title="Tractography Quality Check Report",
    )
    logger.info("HTML report generated: %s", html_output)


def run_all_rounds(
    config: dict,
    results_file: Path,
    output_dir: Path,
    skip_rounds: list[str] | None = None,
    batch_index: int | None = None,
    num_batches: int | None = None,
) -> None:
    """Run all rounds sequentially.
    
    Parameters
    ----------
    config : dict
        Configuration dictionary.
    results_file : Path
        Path to results JSON file.
    output_dir : Path
        Output directory.
    skip_rounds : list[str] | None, optional
        List of rounds to skip.
    batch_index : int | None, optional
        Batch index (1-indexed) if running in batch mode.
    num_batches : int | None, optional
        Total number of batches. Required if batch_index is provided.
    """
    if skip_rounds is None:
        skip_rounds = []

    # Define rounds in order
    rounds = [
        "anatomical",
        "cci",
        "before_after_cci",
        "atlas_views",
        "shape_similarity",
        "afq_profile",
        "bundle_assignment",
    ]

    logger.info("=" * 80)
    logger.info("Running quality check workflow in rounds")
    logger.info("=" * 80)
    logger.info("Config file: %s", config.get("_config_path", "N/A"))
    logger.info("Output directory: %s", output_dir)
    logger.info("Results file: %s", results_file)
    if batch_index is not None:
        logger.info("Batch: %d of %d", batch_index, num_batches or 1)
    logger.info("")

    # Run each round
    for round_name in rounds:
        if round_name in skip_rounds:
            logger.info("Skipping round: %s", round_name)
            continue

        logger.info("")
        logger.info("=" * 80)
        logger.info("Running round: %s", round_name)
        logger.info("=" * 80)

        try:
            run_round(round_name, config, results_file, output_dir, batch_index, num_batches)
            logger.info("Round %s completed successfully", round_name)
        except Exception:
            logger.exception("ERROR: Round %s failed", round_name)
            raise

        # Force garbage collection and memory clearing between rounds
        logger.debug("Clearing memory between rounds...")
        gc.collect()

    # Generate HTML report at the end (only if not in batch mode or if it's the last batch)
    if batch_index is None or (num_batches is not None and batch_index == num_batches):
        logger.info("")
        logger.info("=" * 80)
        logger.info("Generating HTML report")
        logger.info("=" * 80)

        try:
            generate_html(config, results_file, output_dir)
            logger.info("HTML report generated successfully")
        except Exception:
            logger.exception("ERROR: HTML generation failed")
            raise

    logger.info("")
    logger.info("=" * 80)
    logger.info("All rounds completed successfully!")
    logger.info("=" * 80)


def main() -> int:
    """Run the main program."""
    parser = argparse.ArgumentParser(
        description="Run quality check workflow in separate rounds for memory management",
    )
    parser.add_argument(
        "round",
        nargs="?",
        choices=[
            "anatomical",
            "cci",
            "before_after_cci",
            "atlas_views",
            "shape_similarity",
            "afq_profile",
            "bundle_assignment",
            "html",
            "all",
            "generate-slurm",
        ],
        default="all",
        help="Round to run (default: 'all' to run all rounds sequentially). "
        "Use 'generate-slurm' to create a SLURM batch script.",
    )
    parser.add_argument(
        "--config",
        type=Path,
        required=True,
        help="Path to JSON configuration file",
    )
    parser.add_argument(
        "--results-file",
        type=Path,
        default=None,
        help="Path to JSON file for accumulating results (default: <output_dir>/results.json)",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=None,
        help="Output directory (overrides config)",
    )
    parser.add_argument(
        "--skip",
        action="append",
        help="Skip specific rounds (can be used multiple times)",
    )
    parser.add_argument(
        "--verbose",
        "-v",
        action="store_true",
        help="Enable verbose logging",
    )
    parser.add_argument(
        "--batch-index",
        type=int,
        default=None,
        help="Batch index (1-indexed) for SLURM array jobs. "
        "If not provided, will check SLURM_ARRAY_TASK_ID environment variable.",
    )
    parser.add_argument(
        "--num-batches",
        type=int,
        default=None,
        help="Total number of batches. Required when using --batch-index or SLURM arrays.",
    )
    parser.add_argument(
        "--slurm-time",
        type=str,
        default=None,
        help="SLURM time limit (e.g., '24:00:00') for batch script generation",
    )
    parser.add_argument(
        "--slurm-mem",
        type=str,
        default=None,
        help="SLURM memory limit (e.g., '32G') for batch script generation",
    )
    parser.add_argument(
        "--slurm-cpus",
        type=int,
        default=None,
        help="SLURM CPUs per task for batch script generation",
    )
    parser.add_argument(
        "--subjects",
        nargs="+",
        default=None,
        help="List of subject IDs to process (overrides config and batch filtering)",
    )
    parser.add_argument(
        "--subjects-file",
        type=Path,
        default=None,
        help="File containing subject IDs to process (one per line, # for comments)",
    )

    args = parser.parse_args()

    # Setup logging
    log_level = logging.DEBUG if args.verbose else logging.INFO
    logging.basicConfig(
        level=log_level,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )

    # Load config
    config = load_config(args.config)
    # Store config path for logging
    config["_config_path"] = str(args.config)

    # Handle subject filtering from command line
    if args.subjects_file is not None:
        subject_list = load_subject_list(args.subjects_file)
        logger.info("Loaded %d subjects from file: %s", len(subject_list), args.subjects_file)
        config["subject_list"] = subject_list
    elif args.subjects is not None:
        logger.info("Using %d subjects from command line", len(args.subjects))
        config["subject_list"] = args.subjects

    # Determine output directory
    output_dir = args.output_dir
    if output_dir is None:
        output_dir = Path(config.get("output_dir", "output"))
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Determine results file
    results_file = args.results_file
    if results_file is None:
        results_file = output_dir / "results.json"
    results_file = Path(results_file)

    # Determine skip rounds
    skip_rounds = args.skip if args.skip else []

    # Handle batch mode (SLURM array jobs)
    batch_index = args.batch_index
    num_batches = args.num_batches

    # If batch_index not provided, check SLURM_ARRAY_TASK_ID
    if batch_index is None:
        batch_index = get_slurm_array_task_id()
        if batch_index is not None:
            logger.info("Detected SLURM array job with task ID: %d", batch_index)

    # If we have a batch_index, we need num_batches
    if batch_index is not None and num_batches is None:
        logger.error(
            "Batch index provided but num_batches is missing. "
            "Please provide --num-batches or set it in the config file.",
        )
        return 1

    # Check if num_batches is in config
    if num_batches is None:
        num_batches = config.get("num_batches")

    # If we have batch_index, validate it
    if batch_index is not None:
        if num_batches is None:
            logger.error("num_batches must be provided when using batch mode")
            return 1
        if batch_index < 1 or batch_index > num_batches:
            logger.error(
                "Batch index %d is out of range (1-%d)",
                batch_index,
                num_batches,
            )
            return 1

    try:
        if args.round == "generate-slurm":
            # Generate SLURM batch script
            if num_batches is None:
                logger.error("--num-batches is required for generating SLURM batch script")
                return 1

            slurm_options: dict[str, str] = {}
            if args.slurm_time:
                slurm_options["time"] = args.slurm_time
            if args.slurm_mem:
                slurm_options["mem"] = args.slurm_mem
            if args.slurm_cpus:
                slurm_options["cpus-per-task"] = str(args.slurm_cpus)

            rounds_to_run = [r for r in [
                "anatomical",
                "cci",
                "before_after_cci",
                "atlas_views",
                "shape_similarity",
                "afq_profile",
                "bundle_assignment",
            ] if r not in skip_rounds]

            script_path = generate_slurm_batch_script(
                config_file=args.config,
                output_dir=output_dir,
                num_batches=num_batches,
                rounds=rounds_to_run,
                slurm_options=slurm_options,
            )
            logger.info("")
            logger.info("To submit the job array, run:")
            logger.info("  sbatch %s", script_path)
            return 0
        if args.round == "all":
            run_all_rounds(
                config,
                results_file,
                output_dir,
                skip_rounds=skip_rounds,
                batch_index=batch_index,
                num_batches=num_batches,
            )
        elif args.round == "html":
            generate_html(config, results_file, output_dir)
        else:
            run_round(
                args.round,
                config,
                results_file,
                output_dir,
                batch_index=batch_index,
                num_batches=num_batches,
            )
    except Exception:
        logger.exception("Fatal error")
        return 1
    else:
        return 0


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