#!/bin/bash
# HSMR pose-engine install on Apple Silicon (MPS) — Path B (local HSMR/SKEL).
# Clone -> uv venv(3.10) -> deps -> ungated weights -> apply engine_patches/ -> STOP at SKEL gate.
#
#   ./setup_hsmr.sh            full install (re-runnable; STOPS at the SKEL download gate)
#   ./setup_hsmr.sh --patch    re-apply engine_patches/ onto an existing clone and exit
#
# engine_patches/ holds our ARM/headless fixes (torchvision detector drop-in, headless
# pyrender/OpenGL shims, smoke + detector tests). HSMR itself is gitignored under .engine/,
# so a re-clone would drop those fixes — this script restores them. Override the clone
# location with HSMR_DIR=... (used by the durability test against a temp clone).
#
# ponytail: no `set -e` — optional deps (detectron2/chumpy/render) WARN and we still want the
# env + weights down so the smoke test can run and any failure is diagnosable.
set -uo pipefail
REPO="$(cd "$(dirname "$0")" && pwd)"
ENGINE="$REPO/.engine"
HSMR="${HSMR_DIR:-$ENGINE/HSMR}"
PATCHES="$REPO/engine_patches"
LOG="$ENGINE/setup.log"
JOB="$HOME/.jobs/hsmr_setup"

apply_patches(){
  [ -d "$HSMR" ] || { echo "no HSMR clone at $HSMR — clone first"; exit 1; }
  [ -d "$PATCHES" ] || { echo "no engine_patches/ at $PATCHES"; exit 1; }
  ( cd "$PATCHES" && find . -type f ! -name '.DS_Store' ) | sed 's|^\./||' | while read -r f; do
    mkdir -p "$HSMR/$(dirname "$f")"
    cp "$PATCHES/$f" "$HSMR/$f"
    echo "  patched $f"
  done
}

# Standalone re-apply (after a manual HSMR re-clone).
if [ "${1:-}" = "--patch" ]; then
  echo ">>> applying engine_patches/ -> $HSMR"
  apply_patches
  echo "patches applied."
  exit 0
fi

mkdir -p "$ENGINE" "$HOME/.jobs"
exec > >(tee -a "$LOG") 2>&1
hb(){ echo "$(date +%s) $*" > "$JOB"; echo ">>> [$(date +%H:%M:%S)] $*"; }

hb "clone HSMR + submodules (SKEL)"
if [ ! -d "$HSMR/.git" ]; then
  git clone --recurse-submodules https://github.com/IsshikiHugh/HSMR "$HSMR" || { hb "FAIL clone"; exit 1; }
else
  git -C "$HSMR" submodule update --init || true
fi
cd "$HSMR" || { hb "FAIL cd"; exit 1; }

hb "uv venv python 3.10"
uv venv --python 3.10 .venv || { hb "FAIL venv"; exit 1; }
source .venv/bin/activate

hb "install torch + torchvision (MPS wheel)"
uv pip install torch torchvision || { hb "FAIL torch"; exit 1; }

hb "install requirements.txt"
uv pip install -r requirements.txt || hb "WARN requirements had issues"

hb "pin numpy<2 (chumpy/detectron2/smplx need 1.x)"
uv pip install "numpy==1.26.4"

# detectron2/chumpy are NOT needed — engine_patches replaces the detector and the SKEL runtime
# loads pkls without chumpy (see docs/B_ENGINE_SETUP.md). Left out on purpose; they don't build on ARM.

hb "pip install -e HSMR + thirdparty/SKEL"
uv pip install -e . || hb "WARN pip -e . failed"
uv pip install -e thirdparty/SKEL || hb "WARN SKEL editable install failed"

hb "download UNGATED weights from HuggingFace (regressors + HSMR ckpt + ViTPose backbone)"
hf download IsshikiHugh/HSMR-data_inputs \
  --include "body_models/SMPL_to_J19.pkl" "body_models/J_regressor_SKEL_mix_MALE.pkl" "body_models/J_regressor_SMPL_MALE.pkl" "released_models/HSMR-ViTH-r1d1.tar.gz" "backbone/vitpose_backbone.pth" \
  --local-dir data_inputs || hb "WARN hf download issue"
if [ -f data_inputs/released_models/HSMR-ViTH-r1d1.tar.gz ]; then
  hb "extract HSMR checkpoint"
  tar -xzf data_inputs/released_models/HSMR-ViTH-r1d1.tar.gz -C data_inputs/released_models/ && rm -f data_inputs/released_models/HSMR-ViTH-r1d1.tar.gz
fi
mkdir -p data_inputs/body_models/skel

hb "apply engine_patches/ (ARM detector, headless shims, tests)"
apply_patches

hb "DONE-ENV — SKEL gate: John drops skel_models_v1.1/ (v1.1.1) into data_inputs/body_models/skel/"
echo "=== pip freeze (key) ===" ; uv pip freeze | grep -iE 'torch|numpy|smplx|pytorch-lightning|timm' || true