#!/usr/bin/env bash
#SBATCH --job-name=geometry-forcing-causal-control
#SBATCH --array=0-3
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=14
#SBATCH --gres=gpu:1
#SBATCH --time=02:00:00
#SBATCH --account=berzelius-2026-196
#SBATCH --chdir=/proj/cvl/users/x_fahkh2/worldmem-baseline-evals
#SBATCH --output=slurm_logs/%x-%A_%a.out
#SBATCH --error=slurm_logs/%x-%A_%a.out
set -euo pipefail
module load buildenv-gcccuda/12.9.1-gcc11
ROOT=/proj/cvl/users/x_fahkh2/worldmem-baseline-evals
PYTHON=/proj/cvl/users/x_fahkh2/envs/geometry-forcing-eval/bin/python
DATA_ROOT=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/re10k_dfot/real-estate-10k
MODES=(ar8_cfg4 ar12_cfg4 ar12_cfg1 teacher12_cfg4)
MODE=${MODES[$SLURM_ARRAY_TASK_ID]}
OUTPUT=$ROOT/GeometryForcing/outputs/geometry_causal_${MODE}_job${SLURM_ARRAY_JOB_ID}_task${SLURM_ARRAY_TASK_ID}
export TMPDIR=$ROOT/cache
export TRITON_CACHE_DIR=$ROOT/cache
export HF_HOME=$ROOT/cache/huggingface
export TORCH_HOME=$ROOT/cache/torch
export XDG_CACHE_HOME=$ROOT/cache
export MPLCONFIGDIR=$ROOT/cache/matplotlib
export PYTHONPATH=$ROOT/GeometryForcing
export PYTHONNOUSERSITE=1 PYTHONUNBUFFERED=1 OMP_NUM_THREADS=4 MKL_NUM_THREADS=4
export WANDB_MODE=disabled HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1
mkdir -p "$TMPDIR" "$OUTPUT"
"$PYTHON" "$ROOT/deployment/verify_node.py" --expected-gpus 1 --output "$OUTPUT"
cp "$0" "$OUTPUT/launcher.sh"
cp "$ROOT/deployment/geometry_causal_controls.py" "$OUTPUT/probe.py"
srun --ntasks=1 --gpus-per-task=1 --cpus-per-task=14 --cpu-bind=cores \
    "$PYTHON" "$ROOT/deployment/geometry_causal_controls.py" \
    --mode "$MODE" --data-root "$DATA_ROOT" --output "$OUTPUT"