#!/usr/bin/env bash
#SBATCH --job-name=worldmem-baselines-envs
#SBATCH --partition=berzelius-hopper-cpu
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=14
#SBATCH --time=08: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
ENV_ROOT=/proj/cvl/users/x_fahkh2/envs
MINIFORGE_ROOT=$ENV_ROOT/baseline-miniforge3
CONDA=$MINIFORGE_ROOT/bin/conda
MINECRAFT_DATA=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft
RE10K_DATA=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/re10k_dfot/real-estate-10k
OUTPUT=$ROOT/artifacts/environment_setup_$SLURM_JOB_ID
export TMPDIR=$ROOT/cache
export TRITON_CACHE_DIR=$ROOT/cache
export PIP_CACHE_DIR=$ROOT/cache/pip
export CONDA_PKGS_DIRS=$ROOT/cache/conda
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
export PYTHONUNBUFFERED=1
export OMP_NUM_THREADS=4
export MKL_NUM_THREADS=4
export WANDB_MODE=disabled
export CUDA_HOME=$CUDA_ROOT
export MAX_JOBS=2
export NVCC_THREADS=4
export TORCH_CUDA_ARCH_LIST=9.0a
# Link against the toolkit stub on CPU nodes; GPU jobs load the real driver.
export LIBRARY_PATH="$CUDA_HOME/lib64/stubs:${LIBRARY_PATH:-}"
mkdir -p "$TMPDIR" "$PIP_CACHE_DIR" "$CONDA_PKGS_DIRS" "$MPLCONFIGDIR" "$OUTPUT" "$ENV_ROOT"

if [[ ! -x "$CONDA" ]]; then
    INSTALLER=$ROOT/cache/Miniforge3-24.7.1-2-Linux-x86_64.sh
    curl --fail --location --output "$INSTALLER" \
        https://github.com/conda-forge/miniforge/releases/download/24.7.1-2/Miniforge3-24.7.1-2-Linux-x86_64.sh
    bash "$INSTALLER" -b -p "$MINIFORGE_ROOT"
fi
"$CONDA" --version

for ENV_NAME in decmem-eval matrix-game2-eval live-eval geometry-forcing-eval; do
    if [[ ! -x "$ENV_ROOT/$ENV_NAME/bin/python" ]]; then
        "$CONDA" create --prefix "$ENV_ROOT/$ENV_NAME" \
            --override-channels --channel conda-forge python=3.10 pip -y
    fi
done

DECMEM_PYTHON=$ENV_ROOT/decmem-eval/bin/python
"$DECMEM_PYTHON" -m pip install torch==2.5.1 torchvision==0.20.1 --index-url https://download.pytorch.org/whl/cu121
"$DECMEM_PYTHON" -m pip install -r "$ROOT/deployment/requirements_decmem.txt"
# FlashAttention renames its wheel; source and wheel staging must share a filesystem.
PATH="$ENV_ROOT/decmem-eval/bin:$PATH" \
    "$DECMEM_PYTHON" -m pip install flash-attn==2.7.4.post1 --no-build-isolation --no-cache-dir
(
    export PATH="$ENV_ROOT/decmem-eval/bin:$PATH"
    cd "$ROOT/DecMem/video_sparse_attn"
    "$DECMEM_PYTHON" -m pip install --no-build-isolation .
)

MATRIX_PYTHON=$ENV_ROOT/matrix-game2-eval/bin/python
"$MATRIX_PYTHON" -m pip install -r "$ROOT/deployment/requirements_matrix_game2.txt"
# The upstream 2.8.3 wheel needs GLIBC_2.32; build the same version on the target OS.
PATH="$ENV_ROOT/matrix-game2-eval/bin:$PATH" \
    FLASH_ATTENTION_FORCE_BUILD=TRUE FLASH_ATTN_CUDA_ARCHS=90 \
    "$MATRIX_PYTHON" -m pip install flash-attn==2.8.3 \
    --no-build-isolation --no-binary=flash-attn --no-cache-dir --no-deps --force-reinstall

LIVE_PYTHON=$ENV_ROOT/live-eval/bin/python
"$LIVE_PYTHON" -m pip install -r "$ROOT/LIVE/requirements.txt" pytorch-lightning==2.5.2 setuptools==80.9.0

METRIC_PYTHON=$ENV_ROOT/geometry-forcing-eval/bin/python
"$METRIC_PYTHON" -m pip install -r "$ROOT/deployment/requirements_geometry_forcing.txt" -c "$ROOT/deployment/geometry_forcing_constraints.txt"
"$METRIC_PYTHON" -m pip install --target "$ROOT/shared/metric_vendor" --upgrade --no-deps torchmetrics==0.11.4 torch-fidelity==0.3.0

"$METRIC_PYTHON" "$ROOT/deployment/download_assets.py" --root "$ROOT"
"$METRIC_PYTHON" "$ROOT/deployment/cache_metrics.py"
"$METRIC_PYTHON" "$ROOT/deployment/verify_data.py" \
    --minecraft-root "$MINECRAFT_DATA" \
    --re10k-root "$RE10K_DATA" \
    --output "$OUTPUT/data_preflight.json"

for ENV_NAME in decmem-eval matrix-game2-eval live-eval geometry-forcing-eval; do
    "$ENV_ROOT/$ENV_NAME/bin/python" -m pip check
    "$ENV_ROOT/$ENV_NAME/bin/python" -m pip freeze > "$OUTPUT/${ENV_NAME}_environment.txt"
done
printf 'Environment setup completed; environments are under %s and records are under %s\n' "$ENV_ROOT" "$OUTPUT"