#!/bin/bash

# HALPE Environment Setup Script using uv
# This script creates a conda-like environment using uv (much faster than conda)
# and installs all required packages for the propagation-based pruning project

set -e  # Exit on any error

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Function to print colored output
print_status() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Configuration
ENV_NAME="../halpe_uv"
REQUIREMENTS_FILE="requirements_halpe.txt"
PYTHON_VERSION="3.10.0"

print_status "Starting HALPE environment setup with uv..."

# Check if uv is installed
if ! command -v uv &> /dev/null; then
    print_status "uv not found. Installing uv..."
    curl -LsSf https://astral.sh/uv/install.sh | sh
    export PATH="$HOME/.local/bin:$PATH"
    print_success "uv installed successfully"
else
    print_success "uv is already installed"
fi

# Add uv to PATH if not already there
if [[ ":$PATH:" != *":$HOME/.local/bin:"* ]]; then
    export PATH="$HOME/.local/bin:$PATH"
    print_status "Added uv to PATH"
fi

# Check if requirements file exists
if [ ! -f "$REQUIREMENTS_FILE" ]; then
    print_error "Requirements file $REQUIREMENTS_FILE not found!"
    exit 1
fi

print_status "Requirements file found: $REQUIREMENTS_FILE"

# Remove existing environment if it exists
if [ -d "$ENV_NAME" ]; then
    print_warning "Environment $ENV_NAME already exists. Removing it..."
    rm -rf "$ENV_NAME"
    print_success "Removed existing environment"
fi

# Create virtual environment
print_status "Creating virtual environment: $ENV_NAME with Python $PYTHON_VERSION"
uv venv "$ENV_NAME" --python "$PYTHON_VERSION"
print_success "Virtual environment created"

# Activate environment and install packages
print_status "Installing packages from $REQUIREMENTS_FILE..."
print_status "This may take a few minutes..."

# Install packages using uv
uv pip install --python "$ENV_NAME/bin/python" -r "$REQUIREMENTS_FILE"

print_success "All packages installed successfully!"

# Test the environment
print_status "Testing the environment..."

# Activate environment and test imports
source "$ENV_NAME/bin/activate"

# Test key packages
echo "Testing key packages..."
python -c "
import sys
print(f'Python version: {sys.version}')

import torch
print(f'PyTorch: {torch.__version__}')
print(f'CUDA available: {torch.cuda.is_available()}')

import torchvision
print(f'Torchvision: {torchvision.__version__}')

import transformers
print(f'Transformers: {transformers.__version__}')

import huggingface_hub
print(f'HuggingFace Hub: {huggingface_hub.__version__}')

import datasets
print(f'Datasets: {datasets.__version__}')

import peft
print(f'PEFT: {peft.__version__}')

print('All key packages imported successfully!')
"

print_success "Environment setup completed successfully!"

# Create activation script
ACTIVATION_SCRIPT="activate_halpe.sh"
cat > "$ACTIVATION_SCRIPT" << EOF
#!/bin/bash
# HALPE Environment Activation Script
export PATH="\$HOME/.local/bin:\$PATH"
source $ENV_NAME/bin/activate
echo "HALPE environment activated!"
echo "Python: \$(python --version)"
echo "PyTorch: \$(python -c 'import torch; print(torch.__version__)')"
EOF

chmod +x "$ACTIVATION_SCRIPT"

print_success "Created activation script: $ACTIVATION_SCRIPT"

# Print usage instructions
echo ""
echo "=========================================="
echo "HALPE Environment Setup Complete!"
echo "=========================================="
echo ""
echo "To activate the environment:"
echo "  source $ACTIVATION_SCRIPT"
echo "  # or"
echo "  source $ENV_NAME/bin/activate"
echo ""
echo "To deactivate:"
echo "  deactivate"
echo ""
echo "To add uv to your PATH permanently, add this to your ~/.bashrc:"
echo "  export PATH=\"\$HOME/.local/bin:\$PATH\""
echo ""
echo "Environment location: $(realpath $ENV_NAME)"
echo "==========================================" 