#!/bin/bash
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

if ! command -v uv >/dev/null 2>&1; then
    pip install uv
fi
export PATH="$HOME/.local/bin:$PATH"

if [ ! -d "$SCRIPT_DIR/.venv" ]; then
    ( cd "$SCRIPT_DIR" && uv venv --python 3.12 )
    source "$SCRIPT_DIR/.venv/bin/activate"
    uv pip install "lerobot[multi_task_dit,feetech,dataset,training,async] @ git+https://github.com/huggingface/lerobot.git"
else
    source "$SCRIPT_DIR/.venv/bin/activate"
fi

if ! command -v ffmpeg >/dev/null 2>&1; then
    sudo apt update
    sudo apt install -y ffmpeg
fi

# Swap in the dinov3-compatible diffusion files
for d in "$SCRIPT_DIR"/.venv/lib/python*/site-packages/lerobot/policies/diffusion; do
    if [ -d "$d" ]; then
        cp "$SCRIPT_DIR/configuration_diffusion.py" "$d/configuration_diffusion.py"
        cp "$SCRIPT_DIR/modeling_diffusion.py" "$d/modeling_diffusion.py"
    fi
done


python -m lerobot.scripts.lerobot_rollout     --robot.type=so101_follower     --robot.port=/dev/ttyACM0     --robot.id=my_awesome_follower_arm     --robot.cameras="{front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}"     --strategy.type=base     --policy.path="$SCRIPT_DIR"     --policy.num_inference_steps=10     --device=cuda     --compile_warmup_inferences=5     --fps=20     --interpolation_multiplier=3