#!/usr/bin/env bash
#SBATCH --job-name=videovae_inf_tactile
#SBATCH --partition=sharedp
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --gres=gpu:h100:1
#SBATCH --requeue
#SBATCH --exclude=mfmc10
#SBATCH --output=/group2/ct/weihanx/world_model/slurm-logs/videovae_inf_tactile.%j.log
#SBATCH --error=/group2/ct/weihanx/world_model/slurm-logs/videovae_inf_tactile.%j.log

# Inference launcher for finetuned VideoVAE+ tactile checkpoints.
#
# Usage (sbatch): sbatch VideoVAEPlus/scripts/run_inference_tactile.sh [run_name] [data_root] [out_root]
# Usage (local):  ./VideoVAEPlus/scripts/run_inference_tactile.sh [run_name] [data_root] [out_root]
#
# Optional env overrides:
#   CKPT_PATH       Absolute/relative checkpoint path. If unset, auto-picks:
#                   finetune_runs/<run_name>/checkpoints/last.ckpt (fallback: newest *.ckpt)
#   BASE_CONFIG     Inference config template. Default: configs/inference/config_16z_infer_noloss.yaml
#   DEVICE          Inference device, default cuda:0
#   CHUNK_SIZE      Frames per chunk, default 16 (must be divisible by 4)
#   RESOLUTION_H/W  Max decode resolution, defaults 128 x 128
#   CONDA_HOME      Conda install root, default /group2/ct/weihanx/miniconda3
#   CONDA_ENV       Conda env name, default vae
#   PYTHON_BIN      Python binary override

exec 1> >(stdbuf -oL -eL cat) 2>&1
set -euo pipefail
[[ "${DEBUG:-0}" == "1" ]] && set -x

RUN_NAME="${1:-videovae_ft_config_16z_tactile}"
DATA_ROOT="${2:-/group2/ct/weihanx/world_model/VideoVAEPlus/examples/videos/gt}"
OUT_ROOT="${3:-/group2/ct/weihanx/world_model/VideoVAEPlus/examples/videos/recon/${RUN_NAME}}"

REPO="/group2/ct/weihanx/world_model/VideoVAEPlus"
RUN_DIR="$REPO/finetune_runs/$RUN_NAME"
CHECKPOINT_DIR="$RUN_DIR/checkpoints"
BASE_CONFIG="${BASE_CONFIG:-$REPO/configs/inference/config_16z_infer_noloss.yaml}"
CONDA_HOME="${CONDA_HOME:-/group2/ct/weihanx/miniconda3}"
CONDA_ENV="${CONDA_ENV:-vae}"
PYTHON_BIN="${PYTHON_BIN:-$CONDA_HOME/envs/$CONDA_ENV/bin/python}"
DEVICE="${DEVICE:-cuda:0}"
CHUNK_SIZE="${CHUNK_SIZE:-16}"
RESOLUTION_H="${RESOLUTION_H:-128}"
RESOLUTION_W="${RESOLUTION_W:-128}"

mkdir -p /group2/ct/weihanx/world_model/slurm-logs "$OUT_ROOT"

echo "=== VideoVAE+ inference (tactile finetune) ==="
echo "Host:          $(hostname)"
echo "Job ID:        ${SLURM_JOB_ID:-N/A}"
echo "Run name:      $RUN_NAME"
echo "Data root:     $DATA_ROOT"
echo "Out root:      $OUT_ROOT"
echo "Base config:   $BASE_CONFIG"
echo "Python:        $PYTHON_BIN"
echo "Device:        $DEVICE"
echo "Chunk size:    $CHUNK_SIZE"
echo "Resolution:    ${RESOLUTION_H}x${RESOLUTION_W}"
echo

[[ -d "$REPO" ]] || { echo "ERROR: repo not found: $REPO" >&2; exit 2; }
[[ -d "$DATA_ROOT" ]] || { echo "ERROR: data_root not found: $DATA_ROOT" >&2; exit 2; }
[[ -f "$BASE_CONFIG" ]] || { echo "ERROR: base config not found: $BASE_CONFIG" >&2; exit 2; }
[[ -x "$PYTHON_BIN" ]] || { echo "ERROR: python not found: $PYTHON_BIN" >&2; exit 2; }

if [[ -n "${CKPT_PATH:-}" ]]; then
  CKPT_PATH="$CKPT_PATH"
else
  [[ -d "$CHECKPOINT_DIR" ]] || { echo "ERROR: checkpoint dir not found: $CHECKPOINT_DIR" >&2; exit 2; }
  if [[ -f "$CHECKPOINT_DIR/last.ckpt" ]]; then
    CKPT_PATH="$CHECKPOINT_DIR/last.ckpt"
  else
    shopt -s nullglob
    ckpts=("$CHECKPOINT_DIR"/*.ckpt)
    shopt -u nullglob
    [[ ${#ckpts[@]} -gt 0 ]] || { echo "ERROR: no checkpoints found in $CHECKPOINT_DIR" >&2; exit 2; }
    CKPT_PATH="$(ls -t "${ckpts[@]}" | head -n 1)"
  fi
fi

[[ -f "$CKPT_PATH" ]] || { echo "ERROR: checkpoint not found: $CKPT_PATH" >&2; exit 2; }

export OMP_NUM_THREADS=8
export PYTHONUNBUFFERED=1
export PYTHONFAULTHANDLER=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

TMP_CONFIG="${TMPDIR:-/tmp}/videovaeplus_infer_${RUN_NAME}_${SLURM_JOB_ID:-$$}.yaml"
trap 'rm -f "$TMP_CONFIG"' EXIT

echo "[$(date +%H:%M:%S)] checkpoint: $CKPT_PATH"
echo "[$(date +%H:%M:%S)] writing temp config: $TMP_CONFIG"
"$PYTHON_BIN" - "$BASE_CONFIG" "$CKPT_PATH" "$TMP_CONFIG" <<'PY'
import sys
from omegaconf import OmegaConf

base_cfg, ckpt_path, out_cfg = sys.argv[1:]
cfg = OmegaConf.load(base_cfg)
cfg.model.params.ckpt_path = ckpt_path
OmegaConf.save(config=cfg, f=out_cfg)
print(f"temp config written: {out_cfg}")
PY

cd "$REPO"
echo "[$(date +%H:%M:%S)] launching inference_video.py..."
"$PYTHON_BIN" inference_video.py \
  --data_root "$DATA_ROOT" \
  --out_root "$OUT_ROOT" \
  --config_path "$TMP_CONFIG" \
  --device "$DEVICE" \
  --chunk_size "$CHUNK_SIZE" \
  --resolution "$RESOLUTION_H" "$RESOLUTION_W"

echo "[$(date +%H:%M:%S)] Finished."