#!/usr/bin/env bash
# Pose a Mixamo armature in a .blend file to match a reference image.
#
# Usage:
#   ./pose.sh <reference.png> <input.blend> <output.blend>
#
# 1) Runs RTMPose3D inference on the image (conda env: rtmpose3d).
# 2) Opens the blend, finds the first mixamorig armature with a mesh child,
#    retargets the keypoints onto it via pure FK direction matching,
#    captures head yaw, lands the character on the floor, saves to
#    <output.blend>, and writes a front-view preview to <output>.png.
#
# Best results: use a head-on full-body reference image. RTMPose3D's depth
# is noisy from angled views and gets the wrong front/back for joints.
#
# Required environment:
#   - conda env "rtmpose3d" with mmpose + mmdet + mmcv
#   - Blender 4.x at $BLENDER (defaults to /opt/blender-4.3.2-linux-x64/blender)
#   - RTMPose3D model weights at ~/.cache/torch/hub/checkpoints/
#   - mmpose checkout with the rtmpose3d project at $MMPOSE_DIR
#     (defaults to ~/dev/playground/mmpose)

set -euo pipefail

if [[ $# -lt 3 ]]; then
    echo "Usage: $0 <reference.png> <input.blend> <output.blend>" >&2
    exit 1
fi

# Resolve paths against the user's pwd BEFORE any cd.
IMAGE="$(realpath "$1")"
INPUT_BLEND="$(realpath "$2")"
OUTPUT_BLEND="$(realpath -m "$3")"
mkdir -p "$(dirname "$OUTPUT_BLEND")"

# Where this script lives - apply_pose.py is its sibling.
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# External tools / data (override with env vars to relocate).
MMPOSE_DIR="${MMPOSE_DIR:-$HOME/dev/playground/mmpose}"
BLENDER="${BLENDER:-/opt/blender-4.3.2-linux-x64/blender}"
DET_CKPT="${DET_CKPT:-$HOME/.cache/torch/hub/checkpoints/rtmdet_m_8xb32-100e_coco-obj365-person-235e8209.pth}"
POSE_CKPT="${POSE_CKPT:-$HOME/.cache/torch/hub/checkpoints/rtmw3d-l_8xb64_cocktail14-384x288-794dbc78_20240626.pth}"

[[ -d "$MMPOSE_DIR/projects/rtmpose3d" ]] || {
    echo "MMPOSE_DIR does not contain projects/rtmpose3d: $MMPOSE_DIR" >&2; exit 1
}
[[ -x "$BLENDER" ]] || { echo "Blender not found at $BLENDER" >&2; exit 1; }
[[ -s "$DET_CKPT" ]] || { echo "Detection checkpoint missing: $DET_CKPT" >&2; exit 1; }
[[ -s "$POSE_CKPT" ]] || { echo "Pose checkpoint missing: $POSE_CKPT" >&2; exit 1; }

WORK_DIR="$(mktemp -d)"
trap 'rm -rf "$WORK_DIR"' EXIT

echo "[1/2] Running RTMPose3D inference on $IMAGE ..."
# shellcheck disable=SC1091
source "$HOME/miniconda3/etc/profile.d/conda.sh"
conda activate rtmpose3d
cd "$MMPOSE_DIR/projects/rtmpose3d"
export PYTHONPATH="$PWD:${PYTHONPATH:-}"
python demo/body3d_img2pose_demo.py \
    demo/rtmdet_m_640-8xb32_coco-person.py "$DET_CKPT" \
    configs/rtmw3d-l_8xb64_cocktail14-384x288.py "$POSE_CKPT" \
    --input "$IMAGE" --output-root "$WORK_DIR" \
    --save-predictions --device cpu

INPUT_BASE="$(basename "${IMAGE%.*}")"
JSON="$WORK_DIR/results_${INPUT_BASE}.json"
[[ -s "$JSON" ]] || { echo "Inference did not produce $JSON" >&2; exit 1; }

PREVIEW_PNG="${OUTPUT_BLEND%.blend}.png"

echo "[2/2] Retargeting onto armature in $INPUT_BLEND ..."
"$BLENDER" --background "$INPUT_BLEND" \
    --python "$SCRIPT_DIR/apply_pose.py" \
    -- "$JSON" "$OUTPUT_BLEND" "$PREVIEW_PNG"

echo "Done: $OUTPUT_BLEND"
echo "Preview: $PREVIEW_PNG"