#!/bin/bash
# PyTorch CUDA Library Setup Script
# This script sets up the necessary library paths for PyTorch CUDA support

echo "Setting up PyTorch CUDA library paths..."

# Check if conda environment is active
if [ -z "$CONDA_PREFIX" ]; then
    echo "❌ Error: No conda environment detected. Please activate a conda environment first."
    echo "Example: conda activate your_env_name"
    exit 1
fi

# Create activation directory
ACTIVATE_DIR="$CONDA_PREFIX/etc/conda/activate.d"
mkdir -p "$ACTIVATE_DIR"

# Create activation script
cat > "$ACTIVATE_DIR/pytorch_cuda_libs.sh" << 'EOF'
#!/bin/bash
# Set up NVIDIA CUDA library paths for PyTorch

# Find all NVIDIA library directories
NVIDIA_LIB_PATH=$(find $CONDA_PREFIX -path "*/nvidia/*/lib" -type d 2>/dev/null | tr '\n' ':')
CUSPARSELT_LIB_PATH=$(find $CONDA_PREFIX -path "*/cusparselt/lib" -type d 2>/dev/null | tr '\n' ':')

# Add to LD_LIBRARY_PATH
export LD_LIBRARY_PATH="${NVIDIA_LIB_PATH}${CUSPARSELT_LIB_PATH}${LD_LIBRARY_PATH}"
EOF

# Make script executable
chmod +x "$ACTIVATE_DIR/pytorch_cuda_libs.sh"

echo "✅ CUDA library setup complete!"
echo "📁 Activation script created at: $ACTIVATE_DIR/pytorch_cuda_libs.sh"
echo ""
echo "🔄 To apply changes, reactivate your conda environment:"
echo "   conda deactivate && conda activate $(basename $CONDA_PREFIX)"
echo ""
echo "🧪 Test PyTorch CUDA support:"
echo "   python -c \"import torch; print(f'PyTorch {torch.__version__} - CUDA: {torch.cuda.is_available()}')\""