#!/bin/bash
# ============================================================================
# TD (Time Dilation) — One-Command Setup
# ============================================================================
#
# Run this ONCE on a fresh machine with a GPU:
#   chmod +x install.sh && ./install.sh
#
# What it does:
#   1. Installs all Python dependencies
#   2. Downloads the base model (Qwen3-VL-8B-Instruct)
#   3. Downloads the Transport and Merge code
#   4. Sets up output directories
#   5. Verifies GPU access
#   6. Compiles the starter TD file to make sure everything works
#
# After this, just run:
#   python -m td_lang run td_start.td
#
# Requirements:
#   - Python 3.10+
#   - NVIDIA GPU with 24GB+ VRAM (RTX 4090 or better)
#   - ~50GB disk space (models + checkpoints)
#   - Internet connection (first run only)
# ============================================================================

set -e  # Stop on any error

echo "============================================================"
echo "  TD (Time Dilation) — Setup Script"
echo "============================================================"
echo ""

# ── Step 1: Check Python ──
echo "[1/7] Checking Python..."
if ! command -v python3 &> /dev/null; then
    echo "ERROR: Python 3 not found. Install Python 3.10+ first."
    exit 1
fi
PYTHON_VER=$(python3 -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')")
echo "  Python $PYTHON_VER found."

# ── Step 2: Check GPU ──
echo ""
echo "[2/7] Checking GPU..."
if command -v nvidia-smi &> /dev/null; then
    GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)
    GPU_MEM=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader | head -1)
    echo "  GPU: $GPU_NAME ($GPU_MEM)"
else
    echo "  WARNING: nvidia-smi not found. GPU might not be available."
    echo "  Continuing anyway (some features won't work without GPU)."
fi

# ── Step 3: Install Python packages ──
echo ""
echo "[3/7] Installing Python packages..."
echo "  This takes 5-10 minutes on first run."
pip install --break-system-packages -q \
    torch \
    transformers \
    accelerate \
    bitsandbytes \
    peft \
    trl \
    datasets \
    safetensors \
    sentencepiece \
    protobuf \
    scipy \
    lark \
    duckduckgo-search \
    huggingface_hub \
    2>&1 | tail -5

# Unsloth (optional — speeds up training 2x, but can fail on some systems)
echo "  Trying to install Unsloth (optional speed boost)..."
pip install --break-system-packages -q unsloth 2>/dev/null && echo "  Unsloth installed." || echo "  Unsloth not available (that's fine, PEFT fallback works)."

echo "  Packages installed."

# ── Step 4: Download base model ──
echo ""
echo "[4/7] Downloading base model (Qwen3-VL-8B-Instruct)..."
echo "  This is ~16GB. Go grab a coffee."
python3 -c "
from huggingface_hub import snapshot_download
print('  Downloading Qwen/Qwen3-VL-8B-Instruct...')
path = snapshot_download('Qwen/Qwen3-VL-8B-Instruct', local_dir='./models/Qwen3-VL-8B-Instruct')
print(f'  Downloaded to: {path}')
"
echo "  Base model ready."

# ── Step 5: Download Transport and Merge code ──
echo ""
echo "[5/7] Downloading Transport and Merge code..."
if [ ! -d "Cross-Architecture-Merging-for-Large-Language-Models" ]; then
    git clone https://github.com/FedML-AI/Cross-Architecture-Merging-for-Large-Language-Models.git
    echo "  T&M code cloned."
else
    echo "  T&M code already exists, skipping."
fi

# ── Step 6: Set up directories ──
echo ""
echo "[6/7] Setting up directories..."
mkdir -p td_lang_outputs/{checkpoints,snapshots,arena_logs,committed}
echo "  Output directories created."

# ── Step 7: Verify everything works ──
echo ""
echo "[7/7] Verifying installation..."

# Check td_lang compiles
python3 -c "
from td_lang.grammar import parse_td_file
from td_lang.compiler import compile_program
import ast

program = parse_td_file('td_start.td')
code = compile_program(program)
ast.parse(code)
print('  td_lang: OK (td_start.td compiles)')
"

# Check GPU access from Python
python3 -c "
import torch
if torch.cuda.is_available():
    gpu = torch.cuda.get_device_name(0)
    mem = torch.cuda.get_device_properties(0).total_mem / 1024**3
    print(f'  PyTorch GPU: {gpu} ({mem:.0f}GB)')
else:
    print('  PyTorch GPU: NOT AVAILABLE (CPU only)')
"

# Check key libraries
python3 -c "
import transformers, peft, trl, bitsandbytes, lark, datasets
print(f'  transformers: {transformers.__version__}')
print(f'  peft: {peft.__version__}')
print(f'  trl: {trl.__version__}')
print('  All libraries: OK')
"

echo ""
echo "============================================================"
echo "  SETUP COMPLETE!"
echo "============================================================"
echo ""
echo "  To start TD, run:"
echo "    python -m td_lang run td_start.td"
echo ""
echo "  To just compile (preview what it'll do):"
echo "    python -m td_lang compile td_start.td"
echo ""
echo "  To check syntax only:"
echo "    python -m td_lang check td_start.td"
echo ""
echo "============================================================"