#!/usr/bin/env bash
#SBATCH --job-name=matrix-game2-diagnostic
#SBATCH --array=0-4
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=14
#SBATCH --gres=gpu:1
#SBATCH --time=01: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/matrix-game2-eval/bin/python
DATA_ROOT=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft
MODES=(wan_oracle full_original full_weighted i2v_original i2v_weighted)
MODE=${MODES[$SLURM_ARRAY_TASK_ID]}
OUTPUT=$ROOT/Matrix-Game/Matrix-Game-2/outputs/diagnostic_${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 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/diagnose_matrix_game2.py" "$OUTPUT/probe.py"
srun --ntasks=1 --gpus-per-task=1 --cpus-per-task=14 --cpu-bind=cores \
    "$PYTHON" "$ROOT/deployment/diagnose_matrix_game2.py" \
    --mode "$MODE" --data-root "$DATA_ROOT" --output "$OUTPUT"