#!/bin/bash
#SBATCH --job-name=dovla_horizon_sweep
#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:30:00
#SBATCH --output=outputs/hpc/logs/%x_%j.out
#SBATCH --error=outputs/hpc/logs/%x_%j.err

set -euo pipefail

# DECISIVE EXPERIMENT: Does action horizon raise the oracle ceiling?
# Generates PickCube CIL at horizon {4, 8, 16, 32}, measures oracle ceiling each.
# Baseline (horizon=4) oracle for PickCube = 37.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"
DEMO="$SCRATCH_ROOT/maniskill_data/demos/PickCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5"
OUT_ROOT="${OUT_ROOT:-$SCRATCH_ROOT/experiments/horizon_sweep_pickcube}"
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_ROOT" "$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"

for H in 4 8 16 32; do
  OUT_DIR="$OUT_ROOT/h${H}"
  echo "=================================================="
  echo "Generating PickCube horizon=$H, 200 groups, K=16"
  echo "=================================================="
  apptainer exec --nv --env "$ENVS" \
    "$SIF" "$PYTHON" scripts/generate_maniskill_lattice.py \
    --demo "$DEMO" \
    --out "$OUT_DIR" \
    --env-id PickCube-v1 \
    --num-groups 200 \
    --k 16 \
    --horizon "$H" \
    --seed 0 \
    --shard-size 1024 \
    --sim-backend physx_cuda:0 \
    --render-backend cpu \
    --state-storage archive
done

echo ""
echo "=================================================="
echo "ORACLE CEILING BY HORIZON"
echo "=================================================="
apptainer exec --nv --env "$ENVS" "$SIF" "$PYTHON" - <<'PY'
import sys; sys.path.insert(0,'.')
from dovla_cil.data.datasets import CILDataset
import os
root=os.path.expandvars("/scratch/$USER/dovla/experiments/horizon_sweep_pickcube")
print(f"{'horizon':>8} {'groups':>7} {'oracle':>8} {'expert':>8} {'mean_reward_spread':>18}")
for H in [4,8,16,32]:
    d=os.path.join(root,f"h{H}")
    try:
        ds=CILDataset(d)
    except Exception as e:
        print(f"{H:>8}  ERROR: {e}"); continue
    n=len(ds.group_ids); orac=0; exp=0; spreads=[]
    for gid in ds.group_ids:
        recs=ds.get_group(gid)
        if any(r.reward.terminal_success for r in recs): orac+=1
        if any(r.candidate_type=='expert' and r.reward.terminal_success for r in recs): exp+=1
        scores=[r.reward.score for r in recs]
        spreads.append(max(scores)-min(scores))
    ms=sum(spreads)/len(spreads) if spreads else 0
    print(f"{H:>8} {n:>7} {orac/n:>8.4f} {exp/n:>8.4f} {ms:>18.4f}")
print()
print("Baseline reference: horizon=4 PickCube oracle in full collection = 0.3740")
PY