#!/bin/bash

# 适用于 Jetson 等 ARM64 平台

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

echo "========================================"
echo "  开始安装 CUDA TensorRT"
echo "========================================"

# ==================== 系统基础工具 ====================
echo "[1/7] 安装系统基础工具和库..."
sudo apt-get update

sudo apt-get install -y \
    v4l-utils \
    python3-pip \
    libsamplerate0-dev \
    libmpg123-dev \
    portaudio19-dev \

# ==================== ROS Humble 核心包 ====================
echo "[2/7] 安装 ROS Humble 核心功能包..."
sudo apt-get install -y \
    ros-humble-camera-info-manager \
    ros-humble-image-transport-plugins \
    ros-humble-ament-cmake-clang-format \
    ros-humble-rosbridge-suite \

# ==================== Python 依赖 (PyPI) ====================
echo "[3/7] 安装 Python PyPI 依赖..."
# 使用清华镜像加速
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple \
    'pydantic<2.0' \
    loguru \
    paho-mqtt \
    sounddevice \
    jsonschema


# ==================== PyTorch ====================
echo "[5/7] 安装 PyTorch (x86 版本)..."
# 注意：使用 Jetson AI Lab 的 JP6 CUDA 12.6 版本
pip3 install torch torchvision \
    --index-url https://download.pytorch.org/whl/cu126

# ==================== 深度学习库 ====================
echo "[6/7] 安装 Ultralytics (YOLO)..."
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple ultralytics

# ==================== NumPy 版本固定 ====================
echo "[7/7] 固定 NumPy 版本..."
pip install numpy==1.26 -i https://pypi.tuna.tsinghua.edu.cn/simple

echo ""
echo "========================================"
echo "  还剩CUDA（12.6） CUDNN TensorRT（10.4.0） 安装！"
echo "========================================"
echo ""
echo "建议执行以下命令验证安装:"
echo "  ros2 pkg list | grep humble"
echo "  python3 -c 'import torch; print(torch.__version__)'"
echo "  python3 -c 'import numpy; print(numpy.__version__)'"