#!/bin/bash
# Isaac Sim 5.1 + Isaac Lab 2.2.1 完整安装脚本
# 使用 venv + Python 3.11.13 + PyTorch 2.7.0

set -e  # 遇到错误立即退出

echo "=========================================="
echo "Isaac Sim 5.1 + Isaac Lab 2.2.1 安装脚本"
echo "=========================================="

# 检查 Python 3.11 是否已安装
if ! command -v python3.11 &> /dev/null; then
    echo "错误: Python 3.11 未安装"
    echo "请先运行: bash install_python311.sh"
    exit 1
fi

echo "✓ Python 3.11 已安装: $(python3.11 --version)"

# 设置环境变量
VENV_DIR="$HOME/.venvs/isaaclab"
ISAACLAB_DIR="$HOME/robot/tool/IsaacLab"

# 1. 创建 venv
echo ""
echo "步骤 1/7: 创建 venv 环境..."
if [ -d "$VENV_DIR" ]; then
    echo "venv 已存在，跳过创建"
else
    python3.11 -m venv "$VENV_DIR"
    echo "✓ venv 创建完成: $VENV_DIR"
fi

# 2. 激活 venv
echo ""
echo "步骤 2/7: 激活 venv 并升级 pip..."
source "$VENV_DIR/bin/activate"
python --version
pip install --upgrade pip

# 3. 安装 PyTorch 2.7.0
echo ""
echo "步骤 3/7: 安装 PyTorch 2.7.0..."
pip install "torch==2.7.0" "torchvision==0.22.0" --index-url https://download.pytorch.org/whl/cu128
python -c "import torch; print(f'✓ PyTorch {torch.__version__} 安装成功')"

# 4. 安装 Isaac Sim 5.1（强制统一版本）
echo ""
echo "步骤 4/7: 安装 Isaac Sim 5.1..."

# 如已安装旧版本 isaacsim，则先卸载，避免多个版本混在一起
if python -m pip show isaacsim &> /dev/null; then
    CUR_VER=$(python -m pip show isaacsim | awk '/Version:/ {print $2}')
    echo "检测到已安装的 isaacsim 版本: $CUR_VER，先卸载..."
    pip uninstall -y isaacsim
fi

echo "开始安装 isaacsim==5.1.0..."
pip install "isaacsim[all,extscache]==5.1.0" --extra-index-url https://pypi.nvidia.com
echo "✓ Isaac Sim 5.1.0 安装成功"

# 5. 检查系统依赖
echo ""
echo "步骤 5/7: 检查系统依赖..."
if ! command -v cmake &> /dev/null || ! command -v g++ &> /dev/null; then
    echo "警告: cmake 或 build-essential 未安装"
    echo "请运行: sudo apt install -y cmake build-essential"
else
    echo "✓ 系统依赖已满足"
fi

# 6. 克隆 Isaac Lab
echo ""
echo "步骤 6/7: 克隆 Isaac Lab 2.2.1..."
if [ -d "$ISAACLAB_DIR" ]; then
    echo "Isaac Lab 目录已存在，切换到 v2.2.1..."
    cd "$ISAACLAB_DIR"
    git fetch
    git checkout v2.2.1
else
    cd "$HOME"
    git clone https://github.com/isaac-sim/IsaacLab.git
    cd "$ISAACLAB_DIR"
    git checkout v2.2.1
fi
echo "✓ Isaac Lab 2.2.1 已就绪"

# 7. 安装 Isaac Lab
echo ""
echo "步骤 7/7: 安装 Isaac Lab 扩展..."
echo "这可能需要几分钟..."
./isaaclab.sh --install

echo ""
echo "=========================================="
echo "安装完成！"
echo "=========================================="
echo ""
echo "使用方法:"
echo "1. 激活环境: source $VENV_DIR/bin/activate"
echo "2. 进入目录: cd $ISAACLAB_DIR"
echo "3. 运行示例: ./isaaclab.sh -p scripts/tutorials/00_sim/create_empty.py"
echo ""