mac-data-ml-setup.sh

#!/bin/bash
#
# mac-data-ml-setup.sh: A script to set up a modern Data & ML development
# environment on macOS from scratch.
#
# This script will:
# 1. Check for and install Homebrew, and configure the shell environment.
# 2. Install core tools: direnv, uv.
# 3. Prompt the user for a workspace directory name.
# 4. Create the workspace directory and a shared virtual environment inside it.
# 5. Configure direnv to auto-activate the environment.
# 6. Install a curated list of essential data science and ML packages using uv.
#

# --- Configuration & Colors ---
# Color definitions for beautiful output
C_RESET='\033[0m'
C_RED='\033[0;31m'
C_GREEN='\033[0;32m'
C_YELLOW='\033[0;33m'
C_BLUE='\033[0;34m'
C_MAGENTA='\033[0;35m'
C_BOLD='\033[1m'

# --- Helper Functions ---

# Prints a formatted header message
function print_header() {
    echo -e "\n${C_BOLD}${C_MAGENTA}=====================================================${C_RESET}"
    echo -e "${C_BOLD}${C_MAGENTA} $1 ${C_RESET}"
    echo -e "${C_BOLD}${C_MAGENTA}=====================================================${C_RESET}"
}

# Prints a success message
function print_success() {
    echo -e "${C_GREEN}✓ $1${C_RESET}"
}

# Prints an error message and exits
function print_error() {
    echo -e "${C_RED}✗ ERROR: $1${C_RESET}" >&2
    exit 1
}

# Prints an informational message
function print_info() {
    echo -e "${C_BLUE}  $1${C_RESET}"
}

# Checks if a command exists
function command_exists() {
    command -v "$1" &> /dev/null
}

# --- Main Logic ---

function main() {
    # --- Step 1: Welcome and System Check ---
    clear
    print_header "macOS Data & ML Environment Setup"
    print_info "This script will configure your Mac for modern development."
    print_info "It will install Homebrew (if needed), essential tools, and set up a workspace."

    # Ensure we are on macOS
    if [[ "$(uname)" != "Darwin" ]]; then
        print_error "This script is designed for macOS only."
    fi
    
    # --- Step 2: Install and Configure Homebrew ---
    print_header "Checking for Homebrew"
    if ! command_exists brew; then
        print_info "Homebrew not found. Installing now..."
        /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"
        
        # Add Homebrew to PATH for the current script session AND permanently for the user
        local brew_path
        if [[ -x "/opt/homebrew/bin/brew" ]]; then # Apple Silicon
            brew_path="/opt/homebrew/bin/brew"
        elif [[ -x "/usr/local/bin/brew" ]]; then # Intel Macs
            brew_path="/usr/local/bin/brew"
        fi

        if [[ -n "$brew_path" ]]; then
            # Set up for the current script session
            eval "$($brew_path shellenv)"
            
            # Permanently add to the user's shell profile for future sessions
            local shell_config_file
            if [[ "$SHELL" == *"zsh"* ]]; then
                shell_config_file="$HOME/.zshrc"
            elif [[ "$SHELL" == *"bash"* ]]; then
                shell_config_file="$HOME/.bash_profile"
            fi

            if [[ -n "$shell_config_file" ]]; then
                print_info "Adding Homebrew to your shell profile ($shell_config_file)..."
                touch "$shell_config_file" # Ensure the file exists
                if ! grep -q "eval \"\$($brew_path shellenv)\"" "$shell_config_file"; then
                    echo -e '\n# Set up Homebrew environment' >> "$shell_config_file"
                    echo "eval \"\$($brew_path shellenv)\"" >> "$shell_config_file"
                    print_success "Homebrew shell environment configured permanently."
                else
                    print_success "Homebrew shell environment already configured."
                fi
            fi
        fi

        if ! command_exists brew; then
            print_error "Homebrew installation failed. Please install it manually and re-run the script."
        fi
        print_success "Homebrew installed successfully."
    else
        print_success "Homebrew is already installed."
        print_info "Updating Homebrew to ensure all formulae are current..."
        brew update
    fi

    # --- Step 3: Install Core Tools with Brew ---
    print_header "Installing Core Tools: direnv, uv, python"
    local tools=("direnv" "uv")
    for tool in "${tools[@]}"; do
        if ! brew list --formula | grep -q "^${tool%%@*}"; then
            print_info "Installing ${tool}..."
            brew install "$tool"
            print_success "${tool} installed."
        else
            print_success "${tool} is already installed."
        fi
    done

    # --- Step 4: Configure direnv ---
    print_header "Configuring direnv for your shell"
    # Detect shell and add hook if not present
    local shell_config_file
    if [[ "$SHELL" == *"zsh"* ]]; then
        shell_config_file="$HOME/.zshrc"
    elif [[ "$SHELL" == *"bash"* ]]; then
        shell_config_file="$HOME/.bash_profile"
    else
        print_error "Unsupported shell: $SHELL. Please configure direnv manually."
    fi
    
    if ! grep -q 'eval "$(direnv hook zsh)"' "$shell_config_file" && ! grep -q 'eval "$(direnv hook bash)"' "$shell_config_file"; then
        print_info "Adding direnv hook to $shell_config_file..."
        # Use the appropriate hook for the shell
        if [[ "$SHELL" == *"zsh"* ]]; then
            echo -e '\n# direnv hook\neval "$(direnv hook zsh)"' >> "$shell_config_file"
        else
            echo -e '\n# direnv hook\neval "$(direnv hook bash)"' >> "$shell_config_file"
        fi
        print_success "direnv hook added."
    else
        print_success "direnv hook is already configured."
    fi

    # --- Step 5: Create Workspace Directory ---
    print_header "Setting up your Workspace"
    local workspace_name
    while true; do
        read -p "$(echo -e ${C_YELLOW}"Enter a name for your main workspace folder (e.g., Workspace, Projects): "${C_RESET})" workspace_name
        if [[ -z "$workspace_name" ]]; then
            echo -e "${C_RED}Folder name cannot be empty. Please try again.${C_RESET}"
        elif [[ -e "$HOME/$workspace_name" && ! -d "$HOME/$workspace_name" ]]; then
            echo -e "${C_RED}A file named '$workspace_name' already exists. Please choose a different name.${C_RESET}"
        else
            break
        fi
    done

    local workspace_path="$HOME/$workspace_name"
    mkdir -p "$workspace_path"
    cd "$workspace_path" || print_error "Could not navigate to $workspace_path"
    print_success "Workspace created at: $workspace_path"

    # --- Step 6: Create Virtual Environment and direnv config ---
    print_info "Creating a shared Python virtual environment with uv..."
    uv venv
    print_success "Virtual environment '.venv' created."

    print_info "Configuring direnv to auto-load this environment..."
    echo "source .venv/bin/activate" > .envrc
    direnv allow .
    print_success "'.envrc' created and allowed. Environment will now auto-load."

    # --- Step 7: Install Python Packages ---
    print_header "Installing Core Data & ML Python Packages"
    print_info "This may take a few minutes..."
    
    # A comprehensive list of essential packages
    local packages=(
        "jupyterlab"
        "notebook"
        "pandas"
        "polars"
        "numpy"
        "scikit-learn"
        "matplotlib"
        "seaborn"
        "plotly"
        "sqlalchemy"
        "psycopg2-binary"
        "fastapi[all]"
        "uvicorn"
        "pydantic"
        "duckdb"
        "dbt-duckdb"
        "ruff"
    )

    # Use uv to install packages into the new environment
    uv pip install "${packages[@]}"

    if [[ $? -eq 0 ]]; then
        print_success "All Python packages installed successfully."
    else
        print_error "Failed to install one or more Python packages."
    fi
    
    # --- Final Instructions ---
    print_header "🎉 Setup Complete! 🎉"
    echo -e "${C_BOLD}Your modern Data & ML environment is ready.${C_RESET}"
    echo -e "\n${C_YELLOW}--- Key Information ---${C_RESET}"
    echo -e "  ${C_BOLD}Workspace Directory:${C_RESET} $workspace_path"
    echo -e "  ${C_BOLD}Python Environment:${C_RESET}  $workspace_path/.venv"
    echo -e "  ${C_BOLD}Activation:${C_RESET}          Managed automatically by ${C_GREEN}direnv${C_RESET}."
    
    echo -e "\n${C_YELLOW}--- Next Steps ---${C_RESET}"
    echo -e "1. ${C_BOLD}Restart your terminal${C_RESET} or run ${C_GREEN}source $shell_config_file${C_RESET} to apply all shell changes."
    echo -e "2. Navigate to your workspace: ${C_GREEN}cd ~/$workspace_name${C_RESET}"
    echo -e "3. You will see a 'direnv: loading .envrc' message. Your environment is now active!"
    echo -e "4. Start Jupyter Lab with: ${C_GREEN}jupyter lab${C_RESET}"
    echo -e "\nHappy coding!\n"
}

# Run the main function
main