#!/usr/bin/env bash
#SBATCH --job-name=matrix-game2-minecraft-full
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=8
#SBATCH --cpus-per-task=14
#SBATCH --gres=gpu:8
#SBATCH --time=48:00:00
#SBATCH --account=berzelius-2026-196
#SBATCH --array=42-43%1
#SBATCH --chdir=/proj/cvl/users/x_fahkh2/worldmem-baseline-evals/Matrix-Game/Matrix-Game-2
#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
ENV_ROOT=/proj/cvl/users/x_fahkh2/envs
REPO=$ROOT/Matrix-Game/Matrix-Game-2
PYTHON=$ENV_ROOT/matrix-game2-eval/bin/python
METRIC_PYTHON=$ENV_ROOT/geometry-forcing-eval/bin/python
WORLD_MEM_ROOT=$ROOT/references/worldmem
DATA_ROOT=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft
MANIFEST=$ROOT/deployment/manifests/minecraft_test_300.csv
FID_MANIFEST=$ROOT/deployment/manifests/minecraft_test_fid50.csv
SEED=$SLURM_ARRAY_TASK_ID
JOB=$SLURM_ARRAY_JOB_ID
RUN=matrix-game2_minecraft_300cases_seed${SEED}_job${JOB}
OUTPUT=$REPO/outputs/$RUN

export TMPDIR=$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 TRITON_CACHE_DIR=$ROOT/cache
export PYTHONNOUSERSITE=1
export PYTHONUNBUFFERED=1
export OMP_NUM_THREADS=4
export MKL_NUM_THREADS=4
export WANDB_MODE=disabled
export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1
mkdir -p "$TMPDIR" "$OUTPUT"
cd "$REPO"

"$PYTHON" "$ROOT/deployment/verify_node.py" --output "$OUTPUT"
cp "$0" "$OUTPUT/launcher.sh"

srun --ntasks=8 --gpus-per-task=1 --cpus-per-task=14 --cpu-bind=cores --kill-on-bad-exit=1 \
    "$PYTHON" "$ROOT/shared/run_and_score.py" \
    --method matrix-game2 \
    --workspace-root "$ROOT" \
    --metric-python "$METRIC_PYTHON" \
    --worldmem-root "$WORLD_MEM_ROOT" \
    --manifest "$MANIFEST" \
    --data-root "$DATA_ROOT" \
    --output "$OUTPUT" \
    --seed "$SEED" \
    --limit 300 \
    --max-generated 500 \
    --world-size 8

srun --exclusive --exact --nodes=1 --ntasks=1 --gpus-per-task=1 --cpus-per-task=14 \
    "$METRIC_PYTHON" "$ROOT/shared/aggregate.py" \
    --output "$OUTPUT" \
    --manifest "$MANIFEST" \
    --dataset minecraft \
    --fid-manifest "$FID_MANIFEST" \
    --seed "$SEED"