#!/bin/bash
#SBATCH --job-name=ctt_ms_rollout
#SBATCH --account=def-yalda
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=4
#SBATCH --mem=24G
#SBATCH --time=01:00:00
#SBATCH --array=0-0
#SBATCH --output=outputs/hpc/logs/%x_%A_%a.out
#SBATCH --error=outputs/hpc/logs/%x_%A_%a.err

set -euo pipefail

PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
TASK_ID="${SLURM_ARRAY_TASK_ID:-0}"
if [[ -z "${CHECKPOINT:-}" ]]; then
  if [[ -n "${CHECKPOINT_TEMPLATE:-}" ]]; then
    CHECKPOINT="${CHECKPOINT_TEMPLATE//\{seed\}/$TASK_ID}"
  else
    echo "Set CHECKPOINT to a CTT model.pt or CHECKPOINT_TEMPLATE with {seed}" >&2
    exit 2
  fi
fi
SOURCE_INDEX="${SOURCE_INDEX:-data/cil_charts/train/index.json}"
TARGET_INDEX="${TARGET_INDEX:-data/cil_charts/val/index.json}"
if [[ -z "${OUT_DIR:-}" ]]; then
  if [[ -n "${OUT_DIR_TEMPLATE:-}" ]]; then
    OUT_DIR="${OUT_DIR_TEMPLATE//\{seed\}/$TASK_ID}"
  else
    OUT_DIR="runs/ctt_residual_rollout_smoke_seed${TASK_ID}"
  fi
fi

SCRATCH_ROOT="/scratch/$USER/dovla"
SIF="${SIF:-$SCRATCH_ROOT/containers/pytorch_2.7.1_cuda12.8.sif}"
PYTHON="${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"

K="${K:-8}"
POOL_SIZE="${POOL_SIZE:-0}"
NEIGHBORS="${NEIGHBORS:-8}"
MAX_TARGET_CHARTS="${MAX_TARGET_CHARTS:-8}"
GROUP_BATCH_SIZE="${GROUP_BATCH_SIZE:-1}"
DEVICE="${DEVICE:-cpu}"
SIM_BACKEND="${SIM_BACKEND:-physx_cpu}"
RENDER_BACKEND="${RENDER_BACKEND:-cpu}"
RESTORE_TOLERANCE="${RESTORE_TOLERANCE:-1e-5}"
DELTA_SCALE="${DELTA_SCALE:-1.0}"
EXECUTION_ACTION_SCALE="${EXECUTION_ACTION_SCALE:-1.0}"
EXECUTION_ACTION_SCALE_VECTOR="${EXECUTION_ACTION_SCALE_VECTOR:-}"
EXECUTION_ACTION_SCALE_VECTOR_FILE="${EXECUTION_ACTION_SCALE_VECTOR_FILE:-}"
EXECUTION_ACTION_TRANSFORM="${EXECUTION_ACTION_TRANSFORM:-identity}"
BOOTSTRAP_SAMPLES="${BOOTSTRAP_SAMPLES:-200}"
INCLUDE_TARGETS_WITHOUT_POSITIVES="${INCLUDE_TARGETS_WITHOUT_POSITIVES:-0}"
EXCLUDE_SELF_SOURCE="${EXCLUDE_SELF_SOURCE:-0}"
SKIP_METRICS="${SKIP_METRICS:-0}"
CLIP_ACTIONS="${CLIP_ACTIONS:-1}"

if [[ -z "$EXECUTION_ACTION_SCALE_VECTOR" && -n "$EXECUTION_ACTION_SCALE_VECTOR_FILE" ]]; then
  EXECUTION_ACTION_SCALE_VECTOR="$(<"$EXECUTION_ACTION_SCALE_VECTOR_FILE")"
fi

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

RUNTIME_DIR="/tmp/$USER/ctt-ms-rollout-$SLURM_JOB_ID-${SLURM_ARRAY_TASK_ID:-0}"
CACHE_DIR="/tmp/$USER/ctt-ms-rollout-mesa-$SLURM_JOB_ID-${SLURM_ARRAY_TASK_ID:-0}"
mkdir -p "$RUNTIME_DIR" "$CACHE_DIR"
chmod 700 "$RUNTIME_DIR"

export OMP_NUM_THREADS=1
export OPENBLAS_NUM_THREADS=1
export MKL_NUM_THREADS=1
export DOVLA_TORCH_THREADS=1

EXTRA_ARGS=()
if [[ "$INCLUDE_TARGETS_WITHOUT_POSITIVES" == "1" ]]; then
  EXTRA_ARGS+=(--include-targets-without-positives)
fi
if [[ "$EXCLUDE_SELF_SOURCE" == "1" ]]; then
  EXTRA_ARGS+=(--exclude-self-source)
fi
if [[ "$SKIP_METRICS" == "1" ]]; then
  EXTRA_ARGS+=(--skip-metrics)
fi
if [[ "$CLIP_ACTIONS" == "0" ]]; then
  EXTRA_ARGS+=(--disable-action-clipping)
fi

apptainer exec \
  --env "LD_LIBRARY_PATH=$CPU_RENDER_LIBS/lib:$NATIVE_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,DOVLA_TORCH_THREADS=1,MPLBACKEND=Agg,PYTHONDONTWRITEBYTECODE=1" \
  -B "$PROJECT_DIR:$PROJECT_DIR" \
  -B "/scratch/$USER:/scratch/$USER" \
  "$SIF" "$PYTHON" scripts/eval_ctt_generated_rollout.py \
    --checkpoint "$CHECKPOINT" \
    --source-index "$SOURCE_INDEX" \
    --target-index "$TARGET_INDEX" \
    --out-dir "$OUT_DIR" \
    --k "$K" \
    --pool-size "$POOL_SIZE" \
    --neighbors "$NEIGHBORS" \
    --max-target-charts "$MAX_TARGET_CHARTS" \
    --group-batch-size "$GROUP_BATCH_SIZE" \
    --device "$DEVICE" \
    --sim-backend "$SIM_BACKEND" \
    --render-backend "$RENDER_BACKEND" \
    --restore-tolerance "$RESTORE_TOLERANCE" \
    --delta-scale "$DELTA_SCALE" \
    --execution-action-scale "$EXECUTION_ACTION_SCALE" \
    --execution-action-scale-vector "$EXECUTION_ACTION_SCALE_VECTOR" \
    --execution-action-transform "$EXECUTION_ACTION_TRANSFORM" \
    --bootstrap-samples "$BOOTSTRAP_SAMPLES" \
    --no-markdown-report \
    "${EXTRA_ARGS[@]}"