#!/usr/bin/env bash
# Time an ONNX model under tract on a connected Android device.
#
#   ./tools/onnx-probe-on-device.sh models/keypoints/xfeat-1024.onnx [1x1x768x1024]
#
# ## Why this exists
#
# Every model the app runs on the CPU — faces, masks, keypoints — runs under
# tract on NEON when it runs on the tablet, and the desktop figure says
# nothing about that: twenty cores of AVX2 against a tablet's big.LITTLE
# cluster is not a ratio anyone should guess (faces.md §9 measured 2×
# on one workload and would have guessed wrong). NFR-MRG-1 puts a per-frame
# budget on keypoint detection *on the tablet*, and this is how the number
# is measured rather than assumed.
#
# Cross-builds `dr-segment`'s `onnx_probe` example — without the embedded
# segmentation model, which the probe does not need and which is 11 MB the
# push would otherwise carry — pushes it with the model to `/data/local/tmp`,
# runs it twice (tract's first run pays for planning) and removes both.
#
# Not part of CI, which has no device attached.
set -euo pipefail

TARGET=aarch64-linux-android
API=26
DEST=/data/local/tmp

MODEL="${1:?usage: onnx-probe-on-device.sh <model.onnx> [NxCxHxW]}"
SHAPE="${2:-1x1x768x1024}"

ndk="${ANDROID_NDK_HOME:-}"
if [[ -z "$ndk" ]]; then
    ndk=$(find "${ANDROID_HOME:-$HOME/Android/Sdk}/ndk" -maxdepth 1 -mindepth 1 -type d 2>/dev/null |
        sort -V | tail -1)
fi
clang="$ndk/toolchains/llvm/prebuilt/linux-x86_64/bin/$TARGET$API-clang"
[[ -x "$clang" ]] || {
    echo "no $TARGET$API-clang in $ndk — set ANDROID_NDK_HOME" >&2
    exit 1
}
rustup target list --installed | grep -qx "$TARGET" || rustup target add "$TARGET"

adb devices | grep -qw device || {
    echo "no device attached — plug the tablet in and enable USB debugging" >&2
    exit 1
}

echo "building onnx_probe for $TARGET…"
CARGO_TARGET_AARCH64_LINUX_ANDROID_LINKER="$clang" \
CC_aarch64_linux_android="$clang" \
    cargo build -q -p dr-segment --example onnx_probe --release --target "$TARGET" \
        --no-default-features --features semantic
binary="${CARGO_TARGET_DIR:-target}/$TARGET/release/examples/onnx_probe"

echo "pushing $(basename "$binary") and $(basename "$MODEL")…"
adb push "$binary" "$DEST/onnx_probe" >/dev/null
adb push "$MODEL" "$DEST/probe.onnx" >/dev/null
adb shell chmod 755 "$DEST/onnx_probe"
adb shell "$DEST/onnx_probe" "$DEST/probe.onnx" "$SHAPE"
adb shell rm -f "$DEST/onnx_probe" "$DEST/probe.onnx"