#!/usr/bin/env bash
#SBATCH --job-name=decmem-minecraft-stm6-topk27-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/DecMem
#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/DecMem
PYTHON=$ENV_ROOT/decmem-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=decmem_stm6_topk27_minecraft_300cases_seed${SEED}_job${JOB}
OUTPUT=$REPO/outputs/$RUN
PREFLIGHT_OUTPUT=$REPO/outputs/${RUN}_preflight

export TMPDIR=$ROOT/cache TRITON_CACHE_DIR=$ROOT/cache
export HF_HOME=$ROOT/cache/huggingface TORCH_HOME=$ROOT/cache/torch
export PIP_CACHE_DIR=$ROOT/cache/pip CUDA_HOME=$CUDA_ROOT
export XDG_CACHE_HOME=$ROOT/cache MPLCONFIGDIR=$ROOT/cache/matplotlib
export PYTHONNOUSERSITE=1 PYTHONUNBUFFERED=1
export PYTHONPATH=$REPO
export OMP_NUM_THREADS=4 MKL_NUM_THREADS=4
export WANDB_MODE=disabled HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1
mkdir -p "$TMPDIR" "$OUTPUT"
cd "$REPO"

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

# Read-budget ablation: STM6/top-k27 keeps all candidate history and block size4.
# Exercise three denoised blocks and the complete scoring path before the full run.
srun --exclusive --exact --nodes=1 --ntasks=1 --gpus-per-task=1 --cpus-per-task=14 --cpu-bind=cores \
    "$PYTHON" "$ROOT/shared/run_and_score.py" \
    --method decmem --decmem-budget \
    --workspace-root "$ROOT" \
    --metric-python "$METRIC_PYTHON" \
    --worldmem-root "$WORLD_MEM_ROOT" \
    --manifest "$MANIFEST" \
    --data-root "$DATA_ROOT" \
    --output "$PREFLIGHT_OUTPUT" \
    --seed "$SEED" \
    --limit 1 --max-generated 36 --rank 0 --world-size 1

srun --exclusive --exact --nodes=1 --ntasks=1 --gpus-per-task=1 --cpus-per-task=14 \
    "$METRIC_PYTHON" "$ROOT/shared/aggregate.py" \
    --output "$PREFLIGHT_OUTPUT" \
    --manifest "$MANIFEST" \
    --dataset minecraft \
    --seed "$SEED" \
    --expected-frames 36 --limit 1

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 decmem --decmem-budget \
    --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"