#!/usr/bin/env bash
# Install MoMask (text → motion BVH) into vendor/momask + venvs/momask.
# Research code: needs numpy<2 + matplotlib 3.6 + two small compat patches.
set -euo pipefail
cd "$(dirname "$0")/.."

[ -d vendor/momask ] || git clone -q https://github.com/EricGuo5513/momask-codes.git vendor/momask
uv venv venvs/momask --python 3.11 2>/dev/null || true
uv pip install --python venvs/momask/bin/python \
  torch "numpy<2" scipy "matplotlib==3.6.3" tqdm einops gdown ftfy regex \
  "git+https://github.com/openai/CLIP.git"

cd vendor/momask

# HumanML3D checkpoints (188MB, Google Drive)
if [ ! -d checkpoints/t2m/t2m_nlayer8_nhead6_ld384_ff1024_cdp0.1_rvq6ns ]; then
  mkdir -p checkpoints/t2m && cd checkpoints/t2m
  ../../../../venvs/momask/bin/gdown 1vXS7SHJBgWPt59wupQ5UUzhFObrnGkQ0
  unzip -q humanml3d_models.zip && rm humanml3d_models.zip
  cd ../..
fi

# Patch 1: numpy.core.umath_tests was removed — np.matmul is the same op
python3 - <<'PYEOF'
import pathlib
SHIM = ("from types import SimpleNamespace as _SN; "
        "import numpy as _np; ut = _SN(matrix_multiply=_np.matmul)")
for f in ("visualization/Animation.py", "visualization/Quaternions.py"):
    p = pathlib.Path(f); s = p.read_text()
    if "umath_tests" in s:
        p.write_text(s.replace("import numpy.core.umath_tests as ut", SHIM))
PYEOF

# Patch 2: matplotlib removed the ax.lines/ax.collections setters
python3 - <<'PYEOF'
import pathlib
p = pathlib.Path("utils/plot_script.py"); s = p.read_text()
s = s.replace("ax.lines = []", "[l.remove() for l in list(ax.lines)]")
s = s.replace("ax.collections = []", "[c.remove() for c in list(ax.collections)]")
p.write_text(s)
PYEOF

# Wrapper: restore np.float/np.int/... aliases the old code expects
cat > mb_gen.py <<'PYEOF'
"""MODELBEAST wrapper: numpy>=1.24 compat aliases, then run gen_t2m.py."""
import numpy as np
for _n, _v in {"float": float, "int": int, "object": object, "bool": bool, "str": str}.items():
    if not hasattr(np, _n):
        setattr(np, _n, _v)
import runpy
runpy.run_path("gen_t2m.py", run_name="__main__")
PYEOF

echo "momask installed"