#!/usr/bin/env bash
# One-time host setup for the IMPACT registration pipeline (impact-reg-konfai).
#
# The stock elastix-impact 1.0.0 binary (vboussot/ImpactElastix) needs
# glibc >= 2.35 and links LibTorch 2.8.0. This host (Ubuntu 20.04, glibc 2.31,
# env torch 2.14) satisfies neither, so this script builds a self-contained
# shim at ~/.cache/konfai/elastix-impact-shim:
#   * glibc 2.35 loader + libs, extracted from the niftyreg:2.1.1-ubuntu image
#     (ubuntu 22.04 base) -- docker is needed ONLY for this one-time copy;
#   * LibTorch 2.8.0 cu128, self-contained (its own CUDA runtime), from
#     download.pytorch.org;
#   * a wrapper bin/elastix that runs the stock binary through the glibc 2.35
#     loader with the LibTorch 2.8.0 library path (ignoring the inherited
#     LD_LIBRARY_PATH, so the env's torch 2.14 cannot leak in).
# The pipeline selects the shim via KONFAI_ELASTIX_DIR
# (src/registration_pipelines.py::_impact_elastix_root).
#
# Run from the repo root, inside the conda env:
#   source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313
#   pip install impact-reg-konfai        # the CLI itself
#   bash scripts/setup_impact_konfai.sh  # this script
#
# First `register` afterwards auto-downloads from HF/GitHub: the preset app
# (VBoussot/ImpactReg), the feature models (VBoussot/impact-torchscript-models)
# and, on hosts that CAN run the stock binary directly, the binary itself.
set -euo pipefail

KONFAI_CACHE="${KONFAI_CACHE:-$HOME/.cache/konfai}"
SHIM="$KONFAI_CACHE/elastix-impact-shim"
STOCK="$KONFAI_CACHE/elastix-impact"
GLIBC_SRC_DIR="${GLIBC_SRC_DIR:-glibc-2.35}"
GLIBC_IMG="${GLIBC_IMG:-niftyreg:2.1.1-ubuntu}"
LT_VER="2.8.0"
LT_DIR="libtorch-${LT_VER}-cu128"
LT_URL="https://download.pytorch.org/libtorch/cu128/libtorch-shared-with-deps-${LT_VER}%2Bcu128.zip"

echo "[1/5] stock elastix-impact binary (package auto-install)"
if [ ! -x "$STOCK/bin/elastix" ]; then
    python - <<'PY'
from pathlib import Path
from impact_reg_konfai.models.elastix_install import (
    ELASTIX_CACHE, get_elastix_bin, install_elastix_impact)
if not get_elastix_bin(ELASTIX_CACHE).exists():
    # flavor auto-detected (cu128 on this host); downloads + extracts the zip
    install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
PY
fi
[ -x "$STOCK/bin/elastix" ] || { echo "stock binary missing: $STOCK/bin/elastix"; exit 1; }

echo "[2/5] glibc 2.35 from $GLIBC_IMG (docker, one-time)"
if [ ! -x "$KONFAI_CACHE/$GLIBC_SRC_DIR/ld-linux-x86-64.so.2" ]; then
    docker run --rm --entrypoint sh -v "$KONFAI_CACHE":/dest "$GLIBC_IMG" -c \
        "rm -rf /dest/$GLIBC_SRC_DIR && mkdir -p /dest/$GLIBC_SRC_DIR && \
         cp -a /lib/x86_64-linux-gnu/. /dest/$GLIBC_SRC_DIR/ && \
         cp -a /usr/lib/x86_64-linux-gnu/libstdc++.so.6* /dest/$GLIBC_SRC_DIR/ 2>/dev/null || true"
fi

echo "[3/5] LibTorch ${LT_VER} cu128 (self-contained)"
if [ ! -f "$SHIM/lib/$LT_DIR/libtorch/lib/libc10.so" ]; then
    mkdir -p "$SHIM/lib"
    curl -L --fail -o "$KONFAI_CACHE/$LT_DIR.zip" "$LT_URL"
    unzip -q "$KONFAI_CACHE/$LT_DIR.zip" -d "$SHIM/lib/$LT_DIR"
    rm "$KONFAI_CACHE/$LT_DIR.zip"
fi
# the with-deps zip carries the CUDA libs under hashed names; expose the public
# soname elastix links directly
LT_LIB="$SHIM/lib/$LT_DIR/libtorch/lib"
CUDART="$(ls "$LT_LIB"/libcudart-*.so.12 | head -1)"
ln -sf "$(basename "$CUDART")" "$LT_LIB/libcudart.so.12"

echo "[4/5] shim layout + wrapper"
mkdir -p "$SHIM/bin"
cp "$STOCK/bin/elastix" "$SHIM/bin/elastix-real"
[ -d "$STOCK/lib" ] && cp -a "$STOCK/lib/." "$SHIM/lib/"
[ -d "$KONFAI_CACHE/$GLIBC_SRC_DIR" ] && cp -a "$KONFAI_CACHE/$GLIBC_SRC_DIR" "$SHIM/lib/$GLIBC_SRC_DIR"
cat > "$SHIM/bin/elastix" <<'WRAP'
#!/bin/sh
# elastix-impact 1.0.0 shim: the host glibc (2.31) is too old for the binary
# (built against glibc 2.35) and the binary expects LibTorch 2.8.0, which the
# conda env's torch (2.14) does not provide. Run it through the bundled
# glibc 2.35 loader with the bundled self-contained LibTorch 2.8.0 (cu128).
ROOT="$(cd "$(dirname "$0")/.." && pwd)"
LT="$ROOT/lib/libtorch-2.8.0-cu128/libtorch/lib"
GLIBC="$ROOT/lib/glibc-2.35"
exec env -u LD_LIBRARY_PATH -u LD_PRELOAD \
  "$GLIBC/ld-linux-x86-64.so.2" \
  --library-path "$LT:$ROOT/lib:$GLIBC" \
  "$ROOT/bin/elastix-real" "$@"
WRAP
chmod +x "$SHIM/bin/elastix"

echo "[5/5] verify"
KONFAI_ELASTIX_DIR="$SHIM" python - <<'PY'
import os
from pathlib import Path
from impact_reg_konfai.models.elastix_install import get_elastix_bin, loader_env, try_elastix
import subprocess
root = Path(os.environ["KONFAI_ELASTIX_DIR"])
try_elastix(root)
r = subprocess.run([str(get_elastix_bin(root).resolve()), "--extended-version"],
                   capture_output=True, text=True, env=loader_env(root))
assert r.returncode == 0, r.stderr
print(r.stdout.splitlines()[0] if r.stdout else "elastix OK")
print("elastix-impact shim OK:", root)
PY

echo
echo "Done. The impact pipeline picks the shim up automatically"
echo "(src/registration_pipelines.py::_impact_elastix_root), e.g.:"
echo "  python scripts/bench_registration.py --pipelines impact --n-cases 2"