#!/bin/bash

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

cd "$SCRIPT_DIR"

EXTERNAL_DIR="$SCRIPT_DIR/external"

VENV_DIR="$SCRIPT_DIR/.venv"

# mujoco 3.3.1 publishes wheels only for cp312/cp313 (see .python-version).

PYTHON_VERSION="3.13"

echo "=== Setting up diffmjx-al umbrella repo ==="

# Create virtual environment with uv, recreating it if it has the wrong interpreter

if [ -d "$VENV_DIR" ]; then

CURRENT_VERSION="$("$VENV_DIR/bin/python" -c 'import sys; print("%d.%d" % sys.version_info[:2])' 2>/dev/null || true)"

if [ "$CURRENT_VERSION" = "$PYTHON_VERSION" ]; then

echo "Virtual environment already exists at $VENV_DIR (Python $CURRENT_VERSION)"

else

echo "Virtual environment at $VENV_DIR has Python ${CURRENT_VERSION:-unknown}, expected $PYTHON_VERSION"

echo "Recreating it..."

rm -rf "$VENV_DIR"

uv venv --python "$PYTHON_VERSION" "$VENV_DIR"

fi

else

echo "Creating virtual environment..."

uv venv --python "$PYTHON_VERSION" "$VENV_DIR"

fi

# Create external directory for cloned repos

mkdir -p "$EXTERNAL_DIR"

# Clone repositories if they don't already exist

if [ -d "$EXTERNAL_DIR/softjax" ]; then

echo "softjax already cloned"

else

echo "Cloning softjax..."

git clone git@github.com:a-paulus/softjax.git "$EXTERNAL_DIR/softjax"

fi

if [ -d "$EXTERNAL_DIR/mujoco" ]; then

echo "mujoco already cloned, checking out diffmjx branch..."

git -C "$EXTERNAL_DIR/mujoco" fetch origin diffmjx

else

echo "Cloning mujoco (branch: diffmjx)..."

git clone -b diffmjx git@github.com:martius-lab/mujoco.git "$EXTERNAL_DIR/mujoco"

fi

if [ -d "$EXTERNAL_DIR/mjx_diffrax" ]; then

echo "mjx_diffrax already cloned"

else

echo "Cloning mjx_diffrax..."

git clone git@github.com:martius-lab/mjx_diffrax.git "$EXTERNAL_DIR/mjx_diffrax"

fi

# Install all dependencies, including CUDA 12 JAX from the `gpu` extra. This has to

# run after the clones, since mujoco-mjx and mjx-diffrax are editable local paths.

echo ""

echo "Installing dependencies (including JAX with CUDA 12 support)..."

uv sync

echo ""

echo "=== Setup complete ==="

echo "Run experiments with:"

echo " uv run experiments/<experiment>/run.py"