#!/bin/bash
#SBATCH --job-name=dovla_ms_debug
#SBATCH --account=def-yalda_gpu
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=4
# Physics uses CUDA; state-mode material creation uses the CPU Vulkan renderer to avoid
# Vulkan/CUDA device-ordinal mismatches on shared four-GPU nodes.
#SBATCH --gres=gpu:nvidia_h100_80gb_hbm3_1g.10gb:1
#SBATCH --mem=24G
#SBATCH --time=00:20: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"
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_DIR="${OUT_DIR:-$PROJECT_DIR/outputs/hpc/maniskill_debug_cil}"
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
export OPENBLAS_NUM_THREADS=1
export MKL_NUM_THREADS=1
export LP_NUM_THREADS=1

apptainer exec --nv \
  --env "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" \
  "$SIF" "$PYTHON" - <<'PY'
import gymnasium as gym
import mani_skill
import os
import torch

print("torch", torch.__version__, "cuda", torch.cuda.is_available(), torch.cuda.get_device_name(0))
print("vulkan_icd", os.environ.get("VK_ICD_FILENAMES"), "cuda_visible", os.environ.get("CUDA_VISIBLE_DEVICES"))
env = gym.make(
    "PickCube-v1",
    num_envs=1,
    obs_mode="state",
    control_mode="pd_ee_delta_pose",
    render_mode=None,
    sim_backend="physx_cuda:0",
    render_backend="cpu",
)
env.reset(seed=7)
state = {
    section: {name: value.clone() for name, value in values.items()}
    for section, values in env.unwrapped.get_state_dict().items()
}
action = torch.zeros((1, 7), dtype=torch.float32, device=env.unwrapped.device)
env.unwrapped.set_state_dict(state)
env.unwrapped.agent.controller.reset()
restored = env.unwrapped.get_state_dict()
max_error = max(
    float(torch.max(torch.abs(state[section][name] - restored[section][name])).cpu())
    for section in state
    for name in state[section]
)
print("state_restore_max_error", max_error)
assert max_error <= 1e-6

env.unwrapped.step(action)
next_state_1 = {
    section: {name: value.clone() for name, value in values.items()}
    for section, values in env.unwrapped.get_state_dict().items()
}
env.unwrapped.set_state_dict(state)
env.unwrapped.agent.controller.reset()
env.unwrapped.step(action)
next_state_2 = env.unwrapped.get_state_dict()
branch_error = max(
    float(torch.max(torch.abs(next_state_1[section][name] - next_state_2[section][name])).cpu())
    for section in next_state_1
    for name in next_state_1[section]
)
print("deterministic_branch_max_error", branch_error)
assert branch_error <= 1e-5
env.close()
PY

apptainer exec --nv \
  --env "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" \
  "$SIF" "$PYTHON" scripts/generate_maniskill_lattice.py \
  --demo "$DEMO" \
  --out "$OUT_DIR" \
  --num-groups 8 \
  --k 4 \
  --horizon 4 \
  --seed 0 \
  --shard-size 32 \
  --sim-backend physx_cuda:0 \
  --render-backend cpu \
  --state-storage archive

apptainer exec --nv \
  --env "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" \
  "$SIF" "$PYTHON" scripts/inspect_shard.py "$OUT_DIR/manifest.json"