#!/usr/bin/env bash
#SBATCH --job-name=matrix-game2-native-control
#SBATCH --array=0-4
#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/matrix-game2-eval/bin/python
DATA_ROOT=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft
MODES=(official_demo demo_actions neutral recorded small_mouse)
MODE=${MODES[$SLURM_ARRAY_TASK_ID]}
OUTPUT=$ROOT/Matrix-Game/Matrix-Game-2/outputs/diagnostic_native_${MODE}_job${SLURM_ARRAY_JOB_ID}_task${SLURM_ARRAY_TASK_ID}
export TMPDIR=$ROOT/cache
export TRITON_CACHE_DIR=$ROOT/cache
export TORCHINDUCTOR_CACHE_DIR=$ROOT/cache/torchinductor
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/Matrix-Game/Matrix-Game-2
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" "$TORCHINDUCTOR_CACHE_DIR" "$OUTPUT"
"$PYTHON" "$ROOT/deployment/verify_node.py" --expected-gpus 1 --output "$OUTPUT"
cp "$0" "$OUTPUT/launcher.sh"
cp "$ROOT/deployment/matrix_native_controls.py" "$OUTPUT/probe.py"
cp "$ROOT/Matrix-Game/Matrix-Game-2/inference.py" "$ROOT/Matrix-Game/Matrix-Game-2/configs/inference_yaml/inference_universal.yaml" "$OUTPUT/"
srun --ntasks=1 --gpus-per-task=1 --cpus-per-task=14 --cpu-bind=cores \
    "$PYTHON" "$ROOT/deployment/matrix_native_controls.py" \
    --mode "$MODE" --data-root "$DATA_ROOT" --output "$OUTPUT"