#!/bin/bash
#SBATCH --job-name=dovla_pullcube_h16
#SBATCH --account=def-yalda_gpu
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=4
#SBATCH --gres=gpu:nvidia_h100_80gb_hbm3_1g.10gb:1
#SBATCH --mem=24G
#SBATCH --time=01:00:00
#SBATCH --output=outputs/hpc/logs/%x_%j.out
#SBATCH --error=outputs/hpc/logs/%x_%j.err

set -euo pipefail

# Fix PullCube: only 373 pre-success states available (not 500)

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"
CA_BUNDLE="$SCRATCH_ROOT/ca-bundle.crt"
VULKAN_ICD="$CPU_RENDER_LIBS/share/vulkan/icd.d/lvp_icd.x86_64.json"
DEMO="$SCRATCH_ROOT/maniskill_multitask_demos/PullCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5"
OUT_DIR="$SCRATCH_ROOT/experiments/six_task_h16_collection/PullCube-v1"
RUNTIME_DIR="/tmp/$USER/dovla-runtime-$SLURM_JOB_ID"
CACHE_DIR="/tmp/$USER/dovla-mesa-$SLURM_JOB_ID"

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

export OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 LP_NUM_THREADS=1

ENVS="LD_LIBRARY_PATH=$CPU_RENDER_LIBS/lib:$NATIVE_LIBS:/.singularity.d/libs,VK_ICD_FILENAMES=$VULKAN_ICD,VK_DRIVER_FILES=$VULKAN_ICD,XDG_RUNTIME_DIR=$RUNTIME_DIR,MESA_SHADER_CACHE_DIR=$CACHE_DIR,LIBGL_ALWAYS_SOFTWARE=1,LP_NUM_THREADS=1,SSL_CERT_FILE=$CA_BUNDLE,REQUESTS_CA_BUNDLE=$CA_BUNDLE,OMP_NUM_THREADS=1,OPENBLAS_NUM_THREADS=1,MKL_NUM_THREADS=1"

echo "=================================================="
echo "Task: PullCube-v1"
echo "Groups: 373 (max available)"
echo "Horizon: 16"
echo "=================================================="

apptainer exec --nv --env "$ENVS" \
  "$SIF" "$PYTHON" scripts/generate_maniskill_lattice.py \
  --demo "$DEMO" \
  --out "$OUT_DIR" \
  --env-id PullCube-v1 \
  --num-groups 373 \
  --k 16 \
  --horizon 16 \
  --seed 0 \
  --shard-size 1024 \
  --sim-backend physx_cuda:0 \
  --render-backend cpu \
  --state-storage archive

echo ""
echo "✅ PullCube-v1 generation complete"