#!/bin/bash
#SBATCH --job-name=dovla_6task_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=08:00:00
#SBATCH --output=outputs/hpc/logs/%x_%A_%a.out
#SBATCH --error=outputs/hpc/logs/%x_%A_%a.err
#SBATCH --array=0-5

set -euo pipefail

# Generate 6-task CIL collection with horizon=16 (vs baseline h=4)
# Expected: oracle ceiling ~90%+ (vs 42.57% @ h=4)
# This enables policy success 50-70%+ (vs 29.67% @ h=4)

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"
OUT_ROOT="${OUT_ROOT:-$SCRATCH_ROOT/experiments/six_task_h16_collection}"
RUNTIME_DIR="/tmp/$USER/dovla-runtime-$SLURM_JOB_ID"
CACHE_DIR="/tmp/$USER/dovla-mesa-$SLURM_JOB_ID"

# Task array
TASKS=(PickCube-v1 PushCube-v1 PullCube-v1 StackCube-v1 LiftPegUpright-v1 PegInsertionSide-v1)
TASK=${TASKS[$SLURM_ARRAY_TASK_ID]}

# Demo paths
declare -A DEMOS
DEMOS[PickCube-v1]="$SCRATCH_ROOT/maniskill_data/demos/PickCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5"
DEMOS[PushCube-v1]="$SCRATCH_ROOT/maniskill_multitask_demos/PushCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5"
DEMOS[PullCube-v1]="$SCRATCH_ROOT/maniskill_multitask_demos/PullCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5"
DEMOS[StackCube-v1]="$SCRATCH_ROOT/maniskill_multitask_demos/StackCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5"
DEMOS[LiftPegUpright-v1]="$SCRATCH_ROOT/maniskill_multitask_demos/LiftPegUpright-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5"
DEMOS[PegInsertionSide-v1]="$SCRATCH_ROOT/maniskill_multitask_demos/PegInsertionSide-v1/rl/trajectory.h5"

DEMO_PATH="${DEMOS[$TASK]}"
OUT_DIR="$OUT_ROOT/$TASK"

# Groups per task
if [[ "$TASK" == "PickCube-v1" ]]; then
  NUM_GROUPS=1000
else
  NUM_GROUPS=500
fi

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: $TASK"
echo "Groups: $NUM_GROUPS"
echo "Horizon: 16 (vs baseline 4)"
echo "Demo: $DEMO_PATH"
echo "Output: $OUT_DIR"
echo "=================================================="

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

echo ""
echo "✅ $TASK generation complete"