#!/usr/bin/env python3
"""
Script to run evaluation jobs in parallel via SLURM sbatch.
"""

import re
import shlex
import subprocess
from datetime import date
from pathlib import Path
from typing import List, Optional

from common import (
    ENSEMBLE_METHODS,
    ENVS,
    METHOD_CONFIGS,
    ROBUST_MAZE_CHECKPOINT_STEPS,
    SERVER_PATHS,
    SVDD_METHODS,
    find_best_model_checkpoint,
    find_newest_timestamp_dir,
    get_checkpoints,
    get_eval_env_name,
    get_env_folder,
    get_robust_maze_strong_checkpoint,
    require_non_plain_maze_eval_env,
)


# Default conda environment
DEFAULT_CONDA_ENV = "ood-stable"
DEFAULT_CONTAINER_IMAGE = "yrc-bench-procgen.sif"
REPO_ROOT = Path(__file__).resolve().parents[1]

# SLURM configuration
SLURM_CONFIG = {
    "qos": "default",
    "gres": "gpu:1",
    "time": "3-00:00:00",
    "mem": "100G",
    "cpus-per-task": "30",
}

# Default evaluation configuration
EVAL_DEFAULTS = {
    "video_episodes_to_collect": 16,
    "num_levels": 5000,
    "video_filter": "all",
    "cp_rolling_average": "none",
    "video_logging_mode": "folder",
    "video_filter_mode": "any",
    "coverage_fraction": 0.05,
    "svdd_prefix": "dummy05",
}

# Default number of ensemble members (excluding weak agent which is added automatically)
DEFAULT_NUM_ENSEMBLE_MEMBERS = 4

EVAL_ENVS = [*ENVS, "coinrun_proxy_fail", "maze_proxy_fail"]

# Some evaluation environments intentionally reuse training artifacts from a
# base environment. `coinrun_proxy_fail` and `maze_proxy_fail` change the
# Procgen reward/termination behavior at evaluation time, but their
# weak/strong checkpoints, SVDD models, ensemble members, and YAML configs are
# still the normal coinrun/maze artifacts. Keep `args.env` for the environment
# passed to eval_afhp.py; use this alias only for filesystem lookup paths.
ARTIFACT_ENVS = {
    "coinrun_proxy_fail": "coinrun",
    "maze_proxy_fail": "maze",
}

DEFAULT_CONTAINER_BINDS = {
    "cluster3": ["/path/to/cluster3:/path/to/cluster3"],
    "cluster1": ["/path/to/cluster1:/path/to/cluster1"],
    "cluster2": ["/path/to/cluster2:/path/to/cluster2"],
}

# Redirect wandb staging + artifact cache off the (often-full) root volume.
WANDB_DATA_DIR_PATH = "/path/to/scratch/wandb-data"
WANDB_CACHE_DIR_PATH = "/path/to/scratch/wandb-cache"


def build_wandb_env_block(server: str) -> str:
    """Shell lines that point wandb staging/cache at the NAS volume.

    Only emitted on cluster1, where the root volume is contended and the NAS path
    above is available. Other servers have their own dedicated scratch.
    """
    if server != "cluster1":
        return ""
    return (
        f"export WANDB_DATA_DIR='{WANDB_DATA_DIR_PATH}'\n"
        f"export WANDB_CACHE_DIR='{WANDB_CACHE_DIR_PATH}'\n"
        'mkdir -p "$WANDB_DATA_DIR" "$WANDB_CACHE_DIR"'
    )


def get_svdd_feature_type(method: str) -> str:
    """Get the SVDD feature type from method name."""
    if method == "svdd-image":
        return "image"
    elif method == "svdd-latent":
        return "latent"
    return ""


def get_svdd_policy_name(env: str, exp_id: int, method: str) -> str:
    """Get the SVDD policy directory name."""
    feature_type = get_svdd_feature_type(method)
    # Format: svdd_{env}_{feature_type}_exp{id}
    return f"svdd_{env}_{feature_type}_exp{exp_id}"


def get_svdd_model_path(
    env: str,
    exp_id: int,
    method: str,
    svdd_base_path: str,
    svdd_prefix: Optional[str] = None,
) -> Optional[str]:
    """Get the full path to the trained SVDD model file, or None if it doesn't exist."""
    policy_name = get_svdd_policy_name(env, exp_id, method)
    model_base_path = Path(svdd_base_path)
    if svdd_prefix is not None:
        model_base_path = model_base_path / svdd_prefix
    model_file = model_base_path / policy_name / "trained.joblib"
    if model_file.exists():
        return str(model_file)
    return None


def get_svdd_expected_model_path(
    env: str,
    exp_id: int,
    method: str,
    svdd_base_path: str,
    svdd_prefix: Optional[str] = None,
) -> Path:
    policy_name = get_svdd_policy_name(env, exp_id, method)
    model_base_path = Path(svdd_base_path)
    if svdd_prefix is not None:
        model_base_path = model_base_path / svdd_prefix
    return model_base_path / policy_name / "trained.joblib"


def get_ensemble_member_paths(
    env: str,
    exp_id: int,
    checkpoint_base_path: str,
    num_members: int = DEFAULT_NUM_ENSEMBLE_MEMBERS,
) -> List[Optional[str]]:
    """Get paths to ensemble member checkpoints.

    Ensemble members are stored at:
    {checkpoint_base}/ensembles/dummy2_ensemble_{env}_exp{id}_m{member_id}/...

    Args:
        env: Environment name (maze, coinrun)
        exp_id: Experiment ID
        checkpoint_base_path: Base path for checkpoints
        num_members: Number of ensemble members to find

    Returns:
        List of paths to ensemble member checkpoints, or None for missing members
    """
    env_folder = get_env_folder(env)
    base_path = Path(checkpoint_base_path) / env_folder / "ensembles"

    member_paths = []
    for member_id in range(num_members):
        member_parent = base_path / f"dummy2_ensemble_{env}_exp{exp_id}_m{member_id}"
        member_ts_dir = find_newest_timestamp_dir(member_parent)
        member_model = (
            find_best_model_checkpoint(member_ts_dir) if member_ts_dir else None
        )

        if member_model:
            member_paths.append(str(member_model))
        else:
            member_paths.append(None)

    return member_paths


def scale_slurm_mem(mem: str, factor: int) -> str:
    """Scale a SLURM memory string like 100G by an integer factor."""
    match = re.fullmatch(r"(\d+)([KMGTP]?)", mem)
    if match is None:
        raise ValueError(f"Unsupported SLURM mem format: {mem}")

    value, unit = match.groups()
    return f"{int(value) * factor}{unit}"


def shell_join(args: List[str]) -> str:
    """Quote shell command arguments."""
    return " ".join(shlex.quote(str(arg)) for arg in args)


def normalize_container_binds(binds: Optional[List[str]] = None) -> List[str]:
    """Allow repeated or comma-separated bind arguments."""
    normalized = []
    for bind_value in binds or []:
        for bind in bind_value.split(","):
            bind = bind.strip()
            if bind:
                normalized.append(bind)
    return normalized


def get_container_binds(
    server: str, extra_binds: Optional[List[str]] = None
) -> List[str]:
    """Return default server binds plus user-provided Apptainer binds."""
    return [
        *DEFAULT_CONTAINER_BINDS.get(server, []),
        *normalize_container_binds(extra_binds),
    ]


def resolve_container_image(container_image: str, repo_dir: Path = REPO_ROOT) -> Path:
    """Resolve relative container image paths against the repo root."""
    image_path = Path(container_image).expanduser()
    if image_path.is_absolute():
        return image_path
    return repo_dir / image_path


def build_apptainer_command(
    inner_command: str,
    *,
    container_image: Path,
    repo_dir: Path = REPO_ROOT,
    container_binds: Optional[List[str]] = None,
    use_nv: bool = True,
) -> str:
    """Wrap a command so it runs inside the Apptainer image."""
    container_command = (
        "if [[ ! -x /opt/venv/bin/python ]]; then "
        "echo 'Container Python not found: /opt/venv/bin/python'; exit 1; "
        "fi; "
        "export VIRTUAL_ENV=/opt/venv; "
        "export PATH=/opt/venv/bin:${PATH}; "
        "hash -r; "
        f"{inner_command}"
    )
    apptainer_args = [
        "apptainer",
        "exec",
        "--bind",
        f"{repo_dir}:/workspace",
    ]
    if use_nv:
        apptainer_args.insert(2, "--nv")
    for bind in container_binds or []:
        apptainer_args.extend(["--bind", bind])
    apptainer_args.extend(
        [
            "--env",
            "PYTHONUNBUFFERED=1",
            "--env",
            "PYTHONPATH=/workspace",
            "--env",
            "MPLCONFIGDIR=/workspace/.container_cache/matplotlib",
            "--env",
            "XDG_CACHE_HOME=/workspace/.container_cache/xdg",
            "--env",
            "WANDB_CACHE_DIR=/workspace/.container_cache/wandb",
            "--pwd",
            "/workspace",
            str(container_image),
            "bash",
            "-lc",
            container_command,
        ]
    )
    return shell_join(apptainer_args)


def build_runtime_command(
    python_cmd: str,
    *,
    execution: str = "conda",
    container_image: Optional[Path] = None,
    repo_dir: Path = REPO_ROOT,
    container_binds: Optional[List[str]] = None,
) -> str:
    """Build the command run by srun for the selected execution backend."""
    if execution == "conda":
        return python_cmd
    if execution == "apptainer":
        if container_image is None:
            raise ValueError("container_image is required for apptainer execution")
        return build_apptainer_command(
            python_cmd,
            container_image=container_image,
            repo_dir=repo_dir,
            container_binds=container_binds,
        )
    raise ValueError(f"Unsupported execution backend: {execution}")


def build_runtime_setup(
    *,
    execution: str,
    conda_env: str,
    container_image: Optional[Path] = None,
    repo_dir: Path = REPO_ROOT,
) -> str:
    """Build shell setup lines for a Slurm job."""
    if execution == "conda":
        return f"""echo "Using conda env: {conda_env}"
eval "$(conda shell.bash hook)"
conda activate {conda_env}"""
    if execution == "apptainer":
        if container_image is None:
            raise ValueError("container_image is required for apptainer execution")
        return f"""echo "Using Apptainer image: {container_image}"
if command -v module >/dev/null 2>&1; then
    module load apptainer >/dev/null 2>&1 || true
fi
if [[ ! -f {shlex.quote(str(container_image))} ]]; then
    echo "Apptainer image not found: {container_image}"
    exit 1
fi
mkdir -p {shlex.quote(str(repo_dir / ".container_cache"))}"""
    raise ValueError(f"Unsupported execution backend: {execution}")


def build_python_command(eval_args: dict) -> str:
    """Build the eval_afhp.py command for one evaluation."""
    require_non_plain_maze_eval_env(str(eval_args["env_name"]))

    python_args = [
        "python eval_afhp.py",
        f"-c {eval_args['config']}",
        f"-n {eval_args['name']}",
        f"-en {eval_args['env_name']}",
        "-defer_to_oracle",
        f"-experiment_group {eval_args['experiment_group']}",
        f"-video_episodes_to_collect {eval_args['video_episodes_to_collect']}",
        f"-num_levels={eval_args['num_levels']}",
        f"-video_filter {eval_args['video_filter']}",
        f"-cp_rolling_average {eval_args['cp_rolling_average']}",
        f"-video_logging_mode={eval_args['video_logging_mode']}",
        f"-video_filter_mode={eval_args['video_filter_mode']}",
        f"-sim {eval_args['sim']}",
        f"-weak {eval_args['weak']}",
        f"-strong {eval_args['strong']}",
        f"-level_seeds_file {eval_args['level_seeds_file']}",
        f"-coverage_fraction {eval_args['coverage_fraction']}",
    ]

    if eval_args.get("calibration_levels") is not None:
        python_args.append(f"-calibration_levels {eval_args['calibration_levels']}")

    # Add wandb project if specified
    if eval_args.get("wandb_project"):
        python_args.append(f"-wandb_project {eval_args['wandb_project']}")

    # Add SVDD-specific arguments if present
    if eval_args.get("cp_feature"):
        python_args.append(f"-cp_feature {eval_args['cp_feature']}")
    if eval_args.get("svdd_model_path"):
        python_args.append(f"-f_n {eval_args['svdd_model_path']}")

    # Add ensemble-specific arguments if present
    if eval_args.get("ensemble_members"):
        python_args.append("-cp_ensemble_members")
        for member_path in eval_args["ensemble_members"]:
            python_args.append(f"    {member_path}")

    return " ".join(python_args)


def build_sbatch_command(
    job_name: str,
    eval_args: dict,
    conda_env: str,
    log_dir: Path,
    qos: str = "default",
    *,
    server: str = "cluster1",
    execution: str = "conda",
    container_image: Optional[Path] = None,
    repo_dir: Path = REPO_ROOT,
    container_binds: Optional[List[str]] = None,
) -> str:
    """Build the sbatch command string."""
    # Override QOS in SLURM config
    slurm_config = SLURM_CONFIG.copy()
    slurm_config["qos"] = qos
    slurm_args = " ".join(f"--{k}={v}" for k, v in slurm_config.items())
    python_cmd = build_python_command(eval_args)
    runtime_cmd = build_runtime_command(
        python_cmd,
        execution=execution,
        container_image=container_image,
        repo_dir=repo_dir,
        container_binds=container_binds,
    )
    runtime_setup = build_runtime_setup(
        execution=execution,
        conda_env=conda_env,
        container_image=container_image,
        repo_dir=repo_dir,
    )

    sbatch_script = f"""#!/bin/bash
#SBATCH --job-name={job_name}
#SBATCH --output={log_dir}/%x_%j.out
#SBATCH --error={log_dir}/%x_%j.err
{chr(10).join(f"#SBATCH --{k}={v}" for k, v in slurm_config.items())}

{build_wandb_env_block(server)}

{runtime_setup}
srun {slurm_args} {runtime_cmd}
"""
    return sbatch_script


def build_packed_sbatch_command(
    job_name: str,
    job_specs: List[dict],
    conda_env: str,
    log_dir: Path,
    qos: str = "default",
    *,
    server: str = "cluster1",
    execution: str = "conda",
    container_image: Optional[Path] = None,
    repo_dir: Path = REPO_ROOT,
    container_binds: Optional[List[str]] = None,
) -> str:
    """Build an sbatch script that runs multiple evals on one allocated GPU."""
    if not job_specs:
        raise ValueError("job_specs must not be empty")

    slurm_config = SLURM_CONFIG.copy()
    slurm_config["qos"] = qos
    slurm_config["nodes"] = "1"
    slurm_config["ntasks"] = str(len(job_specs))
    slurm_config["mem"] = scale_slurm_mem(slurm_config["mem"], len(job_specs))

    srun_prefix = (
        "srun --overlap --ntasks=1 "
        f"--cpus-per-task={SLURM_CONFIG['cpus-per-task']} --gres=gpu:1"
    )
    srun_commands = []
    for spec in job_specs:
        eval_name = spec["job_name"]
        python_cmd = build_python_command(spec["eval_args"])
        runtime_cmd = build_runtime_command(
            python_cmd,
            execution=execution,
            container_image=container_image,
            repo_dir=repo_dir,
            container_binds=container_binds,
        )
        stdout = log_dir / f"{eval_name}_${{SLURM_JOB_ID}}.out"
        stderr = log_dir / f"{eval_name}_${{SLURM_JOB_ID}}.err"
        srun_commands.append(
            f'{srun_prefix} --output="{stdout}" --error="{stderr}" '
            f'{runtime_cmd} &\npids+=("$!")'
        )

    launch_block = "\n\n".join(srun_commands)
    runtime_setup = build_runtime_setup(
        execution=execution,
        conda_env=conda_env,
        container_image=container_image,
        repo_dir=repo_dir,
    )

    sbatch_script = f"""#!/bin/bash
#SBATCH --job-name={job_name}
#SBATCH --output={log_dir}/%x_%j.out
#SBATCH --error={log_dir}/%x_%j.err
{chr(10).join(f"#SBATCH --{k}={v}" for k, v in slurm_config.items())}

{build_wandb_env_block(server)}

{runtime_setup}

pids=()
{launch_block}

status=0
for pid in "${{pids[@]}}"; do
    if ! wait "${{pid}}"; then
        status=1
    fi
done
exit "${{status}}"
"""
    return sbatch_script


def chunk_job_specs(job_specs: List[dict], chunk_size: int) -> List[List[dict]]:
    """Chunk job specs for packed GPU submissions."""
    if chunk_size <= 0:
        raise ValueError("chunk_size must be positive")
    return [
        job_specs[index : index + chunk_size]
        for index in range(0, len(job_specs), chunk_size)
    ]


def build_packed_job_name(job_specs: List[dict]) -> str:
    """Build a compact Slurm job name for a packed eval chunk."""
    if len(job_specs) == 1:
        return job_specs[0]["job_name"]
    first = job_specs[0]
    exp_ids = [str(spec["exp_id"]) for spec in job_specs]
    prefix = f"{first['env']}_{first['method']}"
    robust_checkpoint_key = first.get("robust_checkpoint_key")
    if robust_checkpoint_key is not None:
        prefix = f"{prefix}_{robust_checkpoint_key}"
    return f"{prefix}_exp{'-'.join(exp_ids)}"


def submit_job(
    job_name: str,
    eval_args: dict,
    conda_env: str,
    log_dir: Path,
    qos: str = "default",
    dry_run: bool = False,
    *,
    server: str = "cluster1",
    execution: str = "conda",
    container_image: Optional[Path] = None,
    repo_dir: Path = REPO_ROOT,
    container_binds: Optional[List[str]] = None,
) -> None:
    """Submit a single job via sbatch."""
    sbatch_script = build_sbatch_command(
        job_name,
        eval_args,
        conda_env,
        log_dir,
        qos,
        server=server,
        execution=execution,
        container_image=container_image,
        repo_dir=repo_dir,
        container_binds=container_binds,
    )

    if dry_run:
        print(f"=== Job: {job_name} ===")
        print(sbatch_script)
        print()
        return

    # Ensure logs directory exists
    log_dir.mkdir(parents=True, exist_ok=True)

    # Submit via sbatch
    result = subprocess.run(
        ["sbatch"],
        input=sbatch_script,
        text=True,
        capture_output=True,
    )

    if result.returncode == 0:
        print(f"Submitted {job_name}: {result.stdout.strip()}")
    else:
        print(f"Failed to submit {job_name}: {result.stderr}")


def submit_packed_job(
    job_name: str,
    job_specs: List[dict],
    conda_env: str,
    log_dir: Path,
    qos: str = "default",
    dry_run: bool = False,
    *,
    server: str = "cluster1",
    execution: str = "conda",
    container_image: Optional[Path] = None,
    repo_dir: Path = REPO_ROOT,
    container_binds: Optional[List[str]] = None,
) -> None:
    """Submit one Slurm job that runs multiple evals on one GPU."""
    sbatch_script = build_packed_sbatch_command(
        job_name,
        job_specs,
        conda_env,
        log_dir,
        qos,
        server=server,
        execution=execution,
        container_image=container_image,
        repo_dir=repo_dir,
        container_binds=container_binds,
    )

    if dry_run:
        print(f"=== Packed Job: {job_name} ===")
        print(sbatch_script)
        print()
        return

    log_dir.mkdir(parents=True, exist_ok=True)
    result = subprocess.run(
        ["sbatch"],
        input=sbatch_script,
        text=True,
        capture_output=True,
    )

    if result.returncode == 0:
        packed_names = ", ".join(spec["job_name"] for spec in job_specs)
        print(f"Submitted {job_name}: {result.stdout.strip()} ({packed_names})")
    else:
        print(f"Failed to submit {job_name}: {result.stderr}")


def build_preflight_command(
    conda_env: str,
    env_name: str,
    *,
    execution: str = "conda",
    container_image: Optional[Path] = None,
    repo_dir: Path = REPO_ROOT,
    container_binds: Optional[List[str]] = None,
) -> List[str]:
    """Build the local preflight command for the selected execution backend."""
    preflight_args = [
        "python",
        "-m",
        "scripts.preflight_eval_env",
        "--env",
        env_name,
    ]
    if execution == "conda":
        return ["conda", "run", "-n", conda_env, *preflight_args]
    if execution == "apptainer":
        if container_image is None:
            raise ValueError("container_image is required for apptainer execution")
        # acs/procgen are pip-installed into /opt/venv inside the image, not
        # editable from /workspace/lib, so skip the repo-local import check.
        container_preflight_args = [*preflight_args, "--skip-local-path-check"]
        apptainer_cmd = build_apptainer_command(
            shell_join(container_preflight_args),
            container_image=container_image,
            repo_dir=repo_dir,
            container_binds=container_binds,
            use_nv=False,
        )
        return [
            "bash",
            "-lc",
            "module load apptainer >/dev/null 2>&1 || true; " + apptainer_cmd,
        ]
    raise ValueError(f"Unsupported execution backend: {execution}")


def run_preflight_check(
    conda_env: str,
    env_name: str,
    *,
    show_output: bool,
    execution: str = "conda",
    container_image: Optional[Path] = None,
    repo_dir: Path = REPO_ROOT,
    container_binds: Optional[List[str]] = None,
) -> bool:
    """Run local dependency checks before submitting SLURM jobs."""
    command = build_preflight_command(
        conda_env,
        env_name,
        execution=execution,
        container_image=container_image,
        repo_dir=repo_dir,
        container_binds=container_binds,
    )
    result = subprocess.run(
        command,
        cwd=REPO_ROOT,
        text=True,
        capture_output=True,
    )

    if show_output or result.returncode != 0:
        print("=== Preflight check ===")
        print("$ " + " ".join(command))
        if result.stdout:
            print(result.stdout, end="" if result.stdout.endswith("\n") else "\n")
        if result.stderr:
            print(result.stderr, end="" if result.stderr.endswith("\n") else "\n")

    if result.returncode != 0:
        print(
            f"Preflight failed with exit code {result.returncode}; not submitting jobs."
        )
        return False

    return True


def main():
    import argparse

    parser = argparse.ArgumentParser(description="Run evaluation jobs via SLURM")
    parser.add_argument(
        "--dry-run", action="store_true", help="Print commands without submitting"
    )
    parser.add_argument(
        "--conda-env",
        default=DEFAULT_CONDA_ENV,
        help=f"Conda environment to use (default: {DEFAULT_CONDA_ENV})",
    )
    parser.add_argument(
        "--execution",
        choices=["conda", "apptainer"],
        default="conda",
        help="Execution backend for Slurm jobs (default: conda).",
    )
    parser.add_argument(
        "--use-container",
        action="store_true",
        help="Shortcut for --execution apptainer.",
    )
    parser.add_argument(
        "--container-image",
        default=DEFAULT_CONTAINER_IMAGE,
        help=(
            "Apptainer .sif image path for --execution apptainer. Relative paths "
            f"are resolved from the repo root (default: {DEFAULT_CONTAINER_IMAGE})."
        ),
    )
    parser.add_argument(
        "--container-bind",
        action="append",
        default=[],
        help=(
            "Extra Apptainer bind mount. Can be repeated or comma-separated. "
            "The repo and selected server data root are bound automatically."
        ),
    )
    parser.add_argument(
        "--server",
        choices=sorted(SERVER_PATHS.keys()),
        default="cluster1",
        help="Server to use for paths (default: cluster1)",
    )
    parser.add_argument(
        "--qos",
        choices=["default", "high"],
        default="default",
        help="SLURM QOS to use (default: default)",
    )
    parser.add_argument(
        "--env", required=True, choices=EVAL_ENVS, help="Environment to evaluate"
    )
    parser.add_argument(
        "--method",
        required=True,
        choices=list(METHOD_CONFIGS.keys()),
        help="Evaluation method",
    )
    parser.add_argument("--prefix", required=True, help="Experiment group prefix")
    parser.add_argument(
        "--exp-ids",
        type=int,
        nargs="+",
        default=[0, 1, 2, 3, 4],
        help="Experiment IDs to run (default: 0 1 2 3 4)",
    )
    parser.add_argument(
        "--runs-per-gpu",
        type=int,
        default=1,
        help=(
            "Number of eval runs to launch concurrently in each one-GPU Slurm "
            "allocation (default: 1)."
        ),
    )
    parser.add_argument(
        "--num-levels",
        type=int,
        default=EVAL_DEFAULTS["num_levels"],
        help="Number of levels",
    )
    parser.add_argument(
        "--video-episodes",
        type=int,
        default=EVAL_DEFAULTS["video_episodes_to_collect"],
        help="Video episodes to collect",
    )
    parser.add_argument(
        "--video-filter", default=EVAL_DEFAULTS["video_filter"], help="Video filter"
    )
    parser.add_argument(
        "--cp-rolling-average",
        default=EVAL_DEFAULTS["cp_rolling_average"],
        help="Coordination policy rolling average",
    )
    parser.add_argument(
        "--video-logging-mode",
        default=EVAL_DEFAULTS["video_logging_mode"],
        help="Video logging mode",
    )
    parser.add_argument(
        "--video-filter-mode",
        default=EVAL_DEFAULTS["video_filter_mode"],
        help="Video filter mode",
    )
    parser.add_argument(
        "--coverage-fraction",
        type=float,
        default=EVAL_DEFAULTS["coverage_fraction"],
        help=f"Coverage fraction for threshold sampling (default: {EVAL_DEFAULTS['coverage_fraction']})",
    )
    parser.add_argument(
        "--calibration-levels",
        type=int,
        default=None,
        help=(
            "Number of validation split seeds to use for percentile calibration "
            "during eval. Default uses the full validation split."
        ),
    )
    parser.add_argument(
        "--wandb-project",
        default=None,
        help="Override wandb project name",
    )
    parser.add_argument(
        "--svdd-prefix",
        default=EVAL_DEFAULTS["svdd_prefix"],
        help=(
            "SVDD training prefix under the server svdd_base path "
            f"(default: {EVAL_DEFAULTS['svdd_prefix']})."
        ),
    )
    # Override checkpoints if needed
    parser.add_argument("--sim", help="Override sim weak checkpoint path")
    parser.add_argument("--weak", help="Override weak checkpoint path")
    parser.add_argument("--strong", help="Override strong checkpoint path")
    robust_group = parser.add_mutually_exclusive_group()
    robust_group.add_argument(
        "--robust200",
        "--robust-200",
        action="store_true",
        help=(
            "For maze evals, replace the strong agent with the random-start "
            "strong policy at 200M timesteps."
        ),
    )
    robust_group.add_argument(
        "--robust400",
        "--robust-400",
        action="store_true",
        help=(
            "For maze evals, replace the strong agent with the random-start "
            "strong policy at 400M timesteps."
        ),
    )
    # Ensemble-specific arguments
    parser.add_argument(
        "--num-ensemble-members",
        type=int,
        default=DEFAULT_NUM_ENSEMBLE_MEMBERS,
        help=f"Number of ensemble members (default: {DEFAULT_NUM_ENSEMBLE_MEMBERS})",
    )
    args = parser.parse_args()

    if args.use_container:
        args.execution = "apptainer"
    if args.runs_per_gpu <= 0:
        parser.error("--runs-per-gpu must be a positive integer.")

    # Get server-specific paths
    paths = SERVER_PATHS[args.server]
    checkpoint_base_path = paths["checkpoint_base"]
    seeds_base_path = paths["seeds_base"]
    svdd_base_path = paths["svdd_base"]
    svdd_prefix = args.svdd_prefix
    repo_dir = REPO_ROOT
    container_image = resolve_container_image(args.container_image, repo_dir)
    container_binds = get_container_binds(args.server, args.container_bind)

    # `args.env` is the public experiment key. `eval_env_name` is the Procgen
    # environment to instantiate, and `artifact_env` is the namespace used for
    # existing experiment artifacts on disk.
    artifact_env = ARTIFACT_ENVS.get(args.env, args.env)
    eval_env_name = get_eval_env_name(args.env)
    require_non_plain_maze_eval_env(eval_env_name)
    robust_checkpoint_key = None
    if args.robust200:
        robust_checkpoint_key = "robust200"
    elif args.robust400:
        robust_checkpoint_key = "robust400"

    if robust_checkpoint_key is not None and artifact_env != "maze":
        print(f"Error: --{robust_checkpoint_key} is currently supported only for maze.")
        return 1
    if robust_checkpoint_key is not None and args.strong:
        print(f"Error: pass either --strong or --{robust_checkpoint_key}, not both.")
        return 1
    output_prefix = (
        f"{args.prefix}_{robust_checkpoint_key}"
        if robust_checkpoint_key is not None
        else args.prefix
    )

    # Configs are artifact-scoped: proxy-fail evaluation reuses coinrun configs
    # and overrides the actual env through -en below.
    config_file = METHOD_CONFIGS[args.method]
    config_path = f"configs/eval/{artifact_env}/{config_file}"

    # Validate config file exists
    if not Path(config_path).exists():
        print(f"Error: Config file not found: {config_path}")
        return 1

    print(f"Execution: {args.execution}")
    if args.execution == "conda":
        print(f"Conda env: {args.conda_env}")
    else:
        print(f"Container image: {container_image}")
        print(f"Container binds: {', '.join(container_binds) or '(none)'}")

    # Build log directory path: base / wandb_project / prefix / date
    log_base = Path(paths["log_base"])
    wandb_project = args.wandb_project or "default"
    log_dir = log_base / wandb_project / output_prefix / date.today().isoformat()
    print(f"Log dir: {log_dir}")

    if not run_preflight_check(
        args.conda_env,
        eval_env_name,
        show_output=args.dry_run,
        execution=args.execution,
        container_image=container_image,
        repo_dir=repo_dir,
        container_binds=container_binds,
    ):
        return 1

    if args.dry_run:
        print(f"Server: {args.server}")
        print(f"Config: {config_path}")
        print(f"Environment: {args.env}")
        if eval_env_name != args.env:
            print(f"Procgen env: {eval_env_name}")
        if artifact_env != args.env:
            print(f"Artifact env: {artifact_env}")
        print(f"Method: {args.method}")
        print(f"Prefix: {args.prefix}")
        if output_prefix != args.prefix:
            print(f"Output prefix: {output_prefix}")
        if robust_checkpoint_key is not None:
            print(
                "Robust strong checkpoint: "
                f"{robust_checkpoint_key} "
                f"({ROBUST_MAZE_CHECKPOINT_STEPS[robust_checkpoint_key]} timesteps)"
            )
        if args.method in SVDD_METHODS:
            print(f"SVDD prefix: {svdd_prefix}")
            print(f"SVDD base: {Path(svdd_base_path) / svdd_prefix}")
        print(f"Runs per GPU: {args.runs_per_gpu}")
        if args.execution == "apptainer":
            print(f"Container image: {container_image}")
            print(f"Container binds: {', '.join(container_binds) or '(none)'}")
        print()

    job_specs = []

    # Loop over experiment IDs
    for exp_id in args.exp_ids:
        # Checkpoints and method-specific models are artifact-scoped. Job names
        # and experiment groups below stay env-scoped so outputs are labeled by
        # the actual evaluation environment.
        checkpoints = get_checkpoints(artifact_env, exp_id, checkpoint_base_path)
        if args.sim:
            checkpoints["sim"] = args.sim
        if args.weak:
            checkpoints["weak"] = args.weak
        if args.strong:
            checkpoints["strong"] = args.strong
        elif robust_checkpoint_key is not None:
            checkpoints["strong"] = get_robust_maze_strong_checkpoint(
                exp_id,
                checkpoint_base_path,
                ROBUST_MAZE_CHECKPOINT_STEPS[robust_checkpoint_key],
            )

        # Get level seeds file path
        level_seeds_file = Path(seeds_base_path) / f"{exp_id}.json"

        # Get SVDD-specific settings if needed
        svdd_model_path = None
        cp_feature = None
        if args.method in SVDD_METHODS:
            svdd_model_path = get_svdd_model_path(
                artifact_env, exp_id, args.method, svdd_base_path, svdd_prefix
            )
            cp_feature = "obs" if args.method == "svdd-image" else "hidden"

        # Get ensemble-specific settings if needed
        ensemble_members = None
        if args.method in ENSEMBLE_METHODS:
            ensemble_members = get_ensemble_member_paths(
                artifact_env, exp_id, checkpoint_base_path, args.num_ensemble_members
            )

        # Validate checkpoints and seeds file exist
        missing = False
        for name, path in checkpoints.items():
            if not Path(path).exists():
                print(f"Warning: exp{exp_id} {name} checkpoint not found: {path}")
                missing = True

        if not level_seeds_file.exists():
            print(
                f"Warning: exp{exp_id} level seeds file not found: {level_seeds_file}"
            )
            missing = True

        if args.method in SVDD_METHODS and svdd_model_path is None:
            expected_model_path = get_svdd_expected_model_path(
                artifact_env, exp_id, args.method, svdd_base_path, svdd_prefix
            )
            print(f"Warning: exp{exp_id} SVDD model not found at {expected_model_path}")
            missing = True

        if args.method in ENSEMBLE_METHODS and ensemble_members:
            for i, member_path in enumerate(ensemble_members):
                if member_path is None:
                    env_folder = get_env_folder(artifact_env)
                    print(
                        f"Warning: exp{exp_id} ensemble member m{i} not found at "
                        f"{checkpoint_base_path}/{env_folder}/ensembles/dummy2_ensemble_{artifact_env}_exp{exp_id}_m{i}/"
                    )
                    missing = True

        if missing:
            print(f"Skipping exp{exp_id} due to missing files\n")
            continue

        # Build job name and experiment group
        name_parts = [args.env, args.method]
        if robust_checkpoint_key is not None:
            name_parts.append(robust_checkpoint_key)
        name_parts.append(f"exp{exp_id}")
        job_name = "_".join(name_parts)
        experiment_group = f"{output_prefix}_{args.env}_exp{exp_id}"

        if args.dry_run:
            print(f"=== exp{exp_id} ===")
            print(f"  Job name: {job_name}")
            print(f"  Experiment group: {experiment_group}")
            print(f"  Weak:   {checkpoints['weak']}")
            print(f"  Strong: {checkpoints['strong']}")
            print(f"  Seeds:  {level_seeds_file}")
            if svdd_model_path:
                print(f"  SVDD model: {svdd_model_path}")
                print(f"  Feature type: {cp_feature}")
            if ensemble_members:
                print(f"  Ensemble members ({len(ensemble_members)}):")
                for i, member_path in enumerate(ensemble_members):
                    print(f"    m{i}: {member_path}")
            print()

        eval_args = {
            "config": config_path,
            "name": job_name,
            # build_sbatch_command emits this as `-en`, so eval_afhp.py creates
            # the requested environment even when config/artifact lookup used an
            # alias such as coinrun.
            "env_name": eval_env_name,
            "experiment_group": experiment_group,
            "num_levels": args.num_levels,
            "video_episodes_to_collect": args.video_episodes,
            "video_filter": args.video_filter,
            "cp_rolling_average": args.cp_rolling_average,
            "video_logging_mode": args.video_logging_mode,
            "video_filter_mode": args.video_filter_mode,
            "coverage_fraction": args.coverage_fraction,
            "calibration_levels": args.calibration_levels,
            "wandb_project": args.wandb_project,
            "level_seeds_file": str(level_seeds_file),
            "svdd_model_path": svdd_model_path,
            "cp_feature": cp_feature,
            "ensemble_members": ensemble_members,
            **checkpoints,
        }

        job_specs.append(
            {
                "job_name": job_name,
                "eval_args": eval_args,
                "exp_id": exp_id,
                "env": args.env,
                "method": args.method,
                "robust_checkpoint_key": robust_checkpoint_key,
            }
        )

    if args.runs_per_gpu == 1:
        for spec in job_specs:
            submit_job(
                spec["job_name"],
                spec["eval_args"],
                args.conda_env,
                log_dir,
                args.qos,
                dry_run=args.dry_run,
                server=args.server,
                execution=args.execution,
                container_image=container_image,
                repo_dir=repo_dir,
                container_binds=container_binds,
            )
    else:
        for chunk in chunk_job_specs(job_specs, args.runs_per_gpu):
            packed_job_name = build_packed_job_name(chunk)
            submit_packed_job(
                packed_job_name,
                chunk,
                args.conda_env,
                log_dir,
                args.qos,
                dry_run=args.dry_run,
                server=args.server,
                execution=args.execution,
                container_image=container_image,
                repo_dir=repo_dir,
                container_binds=container_binds,
            )

    return 0


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