#!/bin/bash
#SBATCH --account=<your-slurm-account>
#SBATCH --partition=l40s
#SBATCH --nodes=1
#SBATCH --ntasks=8
#SBATCH --gres=gpu:l40s:1
#SBATCH --time=24:00:00
#SBATCH --job-name=ddpm_compare
#SBATCH --mail-user=<your-email>   # replace before submitting
#SBATCH --output=slurm-ddpm-compare-%j.out
#SBATCH --error=slurm-ddpm-compare-%j.err

# DDPM-2 vs DDPM-6 comparison (same cluster layout as 6-param training — see reference below).
#
# Reference training script (Slurm + module + paths pattern):
#   <DDPM_ROOT>/april_26/ddpm_hi_lh6/scripts/shell/train_conditional_lh6.sh
#
# Submit from anywhere:
#   sbatch <DDPM_ROOT>/Models/scripts/run_ddpm_comparison.sh
#
# Extra CLI args for compare_ddpm_models.py pass through, e.g. LHS off:
#   sbatch <DDPM_ROOT>/Models/scripts/run_ddpm_comparison.sh --skip-lhs-r2
#
# Override output dir (optional):
#   sbatch --export=OUTPUT_DIR=<DDPM_ROOT>/Models/ddpm_comparison_out_ab \
#     <DDPM_ROOT>/Models/scripts/run_ddpm_comparison.sh
#
# Optional: override DDPM-2 train/val for the combined loss plot (default: bundled JSON in Models/scripts/):
#   sbatch --export=SLURM_2PARAM=/path/to/slurm-2param-ddpm.out \
#     <DDPM_ROOT>/Models/scripts/run_ddpm_comparison.sh
#
# Interactive (same module as training script):
#   module load python/miniconda3-py3.12-usr
#   bash .../run_ddpm_comparison.sh --skip-lhs-r2

set -euo pipefail

ROOT="<DDPM_ROOT>/Models"
cd "$ROOT"

module load python/miniconda3-py3.12-usr

OUT="${OUTPUT_DIR:-${ROOT}/ddpm_comparison_out}"

echo "==============================================="
echo "Job ID: ${SLURM_JOB_ID:-local}"
echo "Job Name: ${SLURM_JOB_NAME:-run_ddpm_comparison}"
echo "Node: ${SLURM_NODELIST:-$(hostname)}"
echo "GPU: ${CUDA_VISIBLE_DEVICES:-n/a}"
echo "Starting Time: $(date)"
echo "Comparison output: ${OUT}"
echo "Reference Slurm recipe: april_26/ddpm_hi_lh6/scripts/shell/train_conditional_lh6.sh"
echo "==============================================="

PY_ARGS=(python "$ROOT/scripts/compare_ddpm_models.py" --output-dir "$OUT")
if [[ -n "${SLURM_2PARAM:-}" ]]; then
  PY_ARGS+=(--slurm-2param "${SLURM_2PARAM}")
fi
PY_ARGS+=("$@")

"${PY_ARGS[@]}"

echo "==============================================="
echo "Artifacts -> ${OUT}"
echo "Finished at: $(date)"
echo "==============================================="