setup.sh

#!/usr/bin/env bash
# ============================================================================
# setup.sh — Project 10: De Novo Proteomics Pipeline
# Pointer-Network Peptide Sequencing + Neural ODE Phosphoproteomics
# + Kinase Activity Inference
# ============================================================================
set -euo pipefail

YELLOW='\033[1;33m'
GREEN='\033[1;32m'
RED='\033[1;31m'
NC='\033[0m'

PROJ_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MASTER_DIR="$(cd "$PROJ_DIR/../amrita_master" && pwd)"

info()  { echo -e "${GREEN}[INFO]${NC}  $*"; }
warn()  { echo -e "${YELLOW}[WARN]${NC}  $*"; }
error() { echo -e "${RED}[ERROR]${NC} $*"; exit 1; }

if [[ -f "$MASTER_DIR/master_setup.sh" ]]; then
    source "$MASTER_DIR/master_setup.sh"
else
    warn "amrita_master/master_setup.sh not found. Using standalone fallback setup functions."
    
    ensure_venv() {
        local project_dir="$1"
        local venv_dir="${2:-$project_dir/.venv}"

        if [[ ! -d "$venv_dir" ]]; then
            info "Creating Python virtual environment at $venv_dir"
            python3 -m venv "$venv_dir"
        fi

        # shellcheck disable=SC1090
        source "$venv_dir/bin/activate"
        info "Using $(python --version) from $(which python)"
    }

    install_common_python_stack() {
        local project_requirements="$1"

        if [[ ! -f "$project_requirements" ]]; then
            echo -e "${RED}[ERROR]${NC} requirements file not found: $project_requirements" >&2
            return 1
        fi

        info "Upgrading pip, setuptools, wheel"
        pip install --upgrade pip setuptools wheel -q

        info "Installing project requirements (this may take a few minutes)..."
        pip install -r "$project_requirements"
    }

    ensure_basic_layout() {
        local project_dir="$1"
        mkdir -p "$project_dir/results" "$project_dir/configs" "$project_dir/src" "$project_dir/data"
    }
fi

# ---------- 1. System-level dependencies -----------------------------------
info "Checking system dependencies..."

if [[ "$OSTYPE" == "darwin"* ]]; then
    if ! command -v brew &>/dev/null; then
        error "Homebrew not found. Install from https://brew.sh"
    fi
    # LaTeX engine
    if ! command -v xelatex &>/dev/null; then
        info "Installing MacTeX (XeLaTeX)..."
        brew install --cask mactex-no-gui 2>/dev/null || warn "MacTeX already installed or install failed — install manually if needed"
    else
        info "XeLaTeX found: $(which xelatex)"
    fi
elif [[ "$OSTYPE" == "linux-gnu"* ]]; then
    info "Installing LaTeX engine and system libs..."
    sudo apt-get update -qq
    sudo apt-get install -y -qq texlive-xetex texlive-fonts-recommended \
        texlive-latex-extra latexmk build-essential curl git
fi

# ---------- 2. Python virtual environment + dependencies -------------------
VENV_DIR="$PROJ_DIR/.venv"
ensure_venv "$PROJ_DIR" "$VENV_DIR"
install_common_python_stack "$PROJ_DIR/requirements.txt"

# ---------- 4. VS Code extensions (optional, non-blocking) ----------------
if command -v code &>/dev/null; then
    info "Installing recommended VS Code extensions..."
    EXTENSIONS=(
        ms-python.python
        ms-toolsai.jupyter
        ms-toolsai.datawrangler
        eamodio.gitlens
        humao.rest-client
        ms-vscode.hexeditor
        ms-vscode.live-server
    )
    for ext in "${EXTENSIONS[@]}"; do
        code --install-extension "$ext" --force 2>/dev/null || true
    done
else
    warn "VS Code CLI 'code' not found — skipping extension install."
fi

# ---------- 5. Create directory scaffold -----------------------------------
info "Ensuring directory structure..."
ensure_basic_layout "$PROJ_DIR"

# ---------- 6. Validate installation --------------------------------------
info "Validating critical imports..."
python -c "
import torch, torch_geometric, torchdiffeq
import pyteomics, pymzml
import numpy, scipy, pandas, sklearn
import matplotlib, plotly, optuna
import pylatex, jinja2
print('All imports OK ✓')
print(f'  PyTorch        : {torch.__version__}')
print(f'  PyG            : {torch_geometric.__version__}')
print(f'  torchdiffeq    : OK')
print(f'  pyteomics      : OK')
print(f'  Optuna         : {optuna.__version__}')
" || error "Import validation failed — check logs above."

echo ""
info "============================================"
info "  Setup complete.  Activate with:"
info "    source ${VENV_DIR}/bin/activate"
info "  Then run:"
info "    python main.py"
info "============================================"