#!/bin/bash
#SBATCH --job-name=dovla_ms_render
#SBATCH --account=def-yalda_cpu
#SBATCH --partition=cpubase_bycore_b1
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=8
#SBATCH --mem=24G
#SBATCH --time=02:00:00
#SBATCH --output=outputs/hpc/logs/%x_%j.out
#SBATCH --error=outputs/hpc/logs/%x_%j.err

set -euo pipefail

PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
SCRATCH_ROOT="/scratch/$USER/dovla"
SIF="$SCRATCH_ROOT/containers/pytorch_2.7.1_cuda12.8.sif"
PYTHON="$SCRATCH_ROOT/envs/maniskill/bin/python"
NATIVE_LIBS="$SCRATCH_ROOT/native_libs/lib"
CPU_RENDER_LIBS="$SCRATCH_ROOT/cpu_render_libs"
VULKAN_ICD="$CPU_RENDER_LIBS/share/vulkan/icd.d/lvp_icd.x86_64.json"

DATASET="${DATASET:?Set DATASET to a generated ManiSkill CIL directory}"
IMAGE_QUALITY="${IMAGE_QUALITY:-90}"
SEED="${SEED:-0}"
OVERWRITE="${OVERWRITE:-0}"
RUNTIME_DIR="/tmp/$USER/dovla-render-$SLURM_JOB_ID"
CACHE_DIR="/tmp/$USER/dovla-render-cache-$SLURM_JOB_ID"

module load StdEnv/2023 apptainer/1.4.5
cd "$PROJECT_DIR"
mkdir -p "$RUNTIME_DIR" "$CACHE_DIR"
chmod 700 "$RUNTIME_DIR"

ARGS=(
  --dataset "$DATASET"
  --render-backend cpu
  --image-quality "$IMAGE_QUALITY"
  --seed "$SEED"
)
if [[ "$OVERWRITE" == "1" ]]; then
  ARGS+=(--overwrite)
fi

apptainer exec \
  --env "LD_LIBRARY_PATH=$CPU_RENDER_LIBS/lib:$NATIVE_LIBS,XDG_RUNTIME_DIR=$RUNTIME_DIR,MESA_SHADER_CACHE_DIR=$CACHE_DIR,LIBGL_ALWAYS_SOFTWARE=1,LP_NUM_THREADS=4,VK_ICD_FILENAMES=$VULKAN_ICD,VK_DRIVER_FILES=$VULKAN_ICD,OMP_NUM_THREADS=1,OPENBLAS_NUM_THREADS=1,MKL_NUM_THREADS=1" \
  "$SIF" "$PYTHON" scripts/render_maniskill_observations.py "${ARGS[@]}"

rm -rf "$RUNTIME_DIR" "$CACHE_DIR"