#!/bin/bash
#SBATCH --job-name=eval_h16_rollout
#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=32G
#SBATCH --time=02:00:00
#SBATCH --output=outputs/hpc/logs/%x_%A_%a.out
#SBATCH --error=outputs/hpc/logs/%x_%A_%a.err
#SBATCH --array=0-2

set -euo pipefail

# Evaluate h=16 policy via online ManiSkill rollout
# This is the SOTA-comparable metric

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"
RUNTIME_DIR="/tmp/$USER/dovla-runtime-$SLURM_JOB_ID"
CACHE_DIR="/tmp/$USER/dovla-mesa-$SLURM_JOB_ID"

SEED=$SLURM_ARRAY_TASK_ID
CHECKPOINT="$SCRATCH_ROOT/experiments/dovla_h16_rollout_runs/seed_$SEED/best.pt"
DATASET="$SCRATCH_ROOT/experiments/six_task_h16_collection"
OUT="$SCRATCH_ROOT/experiments/dovla_h16_rollout_runs/seed_${SEED}/online_rollout.json"

module load StdEnv/2023 apptainer/1.4.5
cd "$PROJECT_DIR"
mkdir -p outputs/hpc/logs "$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 "Online Rollout Evaluation - h=16 Policy"
echo "Seed: $SEED"
echo "Checkpoint: $CHECKPOINT"
echo "Dataset: $DATASET"
echo "Expected: 55-70%+ policy success (vs 29.67% baseline)"
echo "=================================================="

apptainer exec --nv --env "$ENVS" \
  "$SIF" "$PYTHON" scripts/eval_maniskill_policy_rollout.py \
  --checkpoint "$CHECKPOINT" \
  --dataset "$DATASET" \
  --out "$OUT" \
  --device cuda \
  --max-groups 700 \
  --group-batch-size 16 \
  --sim-backend physx_cuda:0 \
  --render-backend cpu

echo ""
echo "✅ Rollout evaluation complete for seed $SEED"
echo "Results: $OUT"