#!/bin/bash
set -e  # Exit immediately if a command exits with a non-zero status
# Update system packages
echo "Updating system packages..."
sudo apt update && sudo apt upgrade -y
# Install system dependencies
echo "Installing system dependencies..."
sudo apt install vim ffmpeg libsndfile1 git htop git-lfs -y

# Install Python 3.10
echo "Installing Python 3.10..."
sudo apt install python3.10 python3.10-venv -y

# Check if Python 3.10 is installed
if ! command -v python3.10 &> /dev/null; then
    echo "Failed to install Python 3.10. Please check your system and try again."
    exit 1
fi

# Get Python version
PYTHON_VERSION=$(python3.10 --version | awk '{print $2}')
echo "Using Python version: ${PYTHON_VERSION}"

# Check for CUDA
if command -v nvidia-smi &> /dev/null; then
    echo "NVIDIA GPU detected. Will install CUDA-enabled PyTorch."
    CUDA_AVAILABLE=1
    CUDA_VERSION=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -n 1)
    echo "CUDA Driver Version: ${CUDA_VERSION}"
else
    echo "No NVIDIA GPU detected. Will install CPU-only PyTorch."
    CUDA_AVAILABLE=0
fi

# Create a virtual environment if it doesn't exist
if [ ! -d "venv" ]; then
    echo "Creating Python virtual environment..."
    python3.10 -m venv venv
else
    echo "Virtual environment already exists."
fi

# Activate virtual environment
source venv/bin/activate

# Upgrade pip
echo "Upgrading pip..."
pip install --upgrade pip

# Get CUDA version number if available
if [ "$CUDA_AVAILABLE" -eq 1 ]; then
    CUDA_VERSION_NUM=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -n 1 | cut -d'.' -f1)
    # Map CUDA driver version to PyTorch CUDA version
    if [ "$CUDA_VERSION_NUM" -ge 535 ]; then
        PYTORCH_CUDA="cu124"  # CUDA 12.4
    elif [ "$CUDA_VERSION_NUM" -ge 530 ]; then
        PYTORCH_CUDA="cu121"  # CUDA 12.1
    elif [ "$CUDA_VERSION_NUM" -ge 520 ]; then
        PYTORCH_CUDA="cu118"  # CUDA 11.8
    elif [ "$CUDA_VERSION_NUM" -ge 510 ]; then
        PYTORCH_CUDA="cu117"  # CUDA 11.7
    else
        echo "Warning: CUDA version may be too old. Defaulting to CPU version."
        CUDA_AVAILABLE=0
    fi
fi

# Install PyTorch 2.6.0
echo "Installing PyTorch 2.6.0..."
if [ "$CUDA_AVAILABLE" -eq 1 ]; then
    echo "Installing PyTorch with CUDA support (${PYTORCH_CUDA})..."
    pip install torch==2.6.0 --index-url https://download.pytorch.org/whl/${PYTORCH_CUDA}
else
    echo "Installing CPU-only PyTorch..."
    pip install torch==2.6.0
fi

# Install Hugging Face CLI and tools
echo "Installing Hugging Face CLI and tools..."
pip install huggingface_hub

# Verify installation
echo "Verifying installation..."
python -c "import torch; print('PyTorch version:', torch.__version__); print('CUDA available:', torch.cuda.is_available())"

# Check for Hugging Face CLI
python -c "import huggingface_hub; print('Hugging Face Hub version:', huggingface_hub.__version__)"

echo "============================================================"
echo "Setup complete!"
echo "To activate this environment, run:"
echo "source venv/bin/activate"
echo "============================================================"