#!/bin/bash
# ==========================================================================
# ArcisVLM Full Training Pipeline — 8x A100 DDP
#
# Runs all stages sequentially:
#   Stage 0: Train 32K BPE tokenizer
#   Stage 1: JEPA Alignment (~2.5 hrs)
#   Stage 2: Instruction Tuning (~1.5 hrs)
#   Stage 3: Domain Fine-Tuning (~30 min)
#
# Usage:
#   bash scripts/launch_training.sh
# ==========================================================================
set -e

CONFIG=configs/scale_1.3b.yaml
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
export OMP_NUM_THREADS=4
export NCCL_IB_DISABLE=0
export NCCL_P2P_LEVEL=NVL

# Change to project root
cd "$(dirname "$0")/.."

echo "=================================================================="
echo "ArcisVLM Training Pipeline"
echo "=================================================================="
echo "Config:     $CONFIG"
echo "GPUs:       $CUDA_VISIBLE_DEVICES"
echo "Start time: $(date)"
echo "=================================================================="

# ------------------------------------------------------------------
echo ""
echo "=== STAGE 0: TRAIN TOKENIZER ==="
echo "Started: $(date)"
python3 scripts/train_tokenizer_32k.py
echo "Completed: $(date)"

# ------------------------------------------------------------------
echo ""
echo "=== STAGE 1: JEPA ALIGNMENT (~2.5 hrs) ==="
echo "Started: $(date)"
torchrun --nproc_per_node=8 scripts/train_stage1_ddp.py --config $CONFIG
echo "Completed: $(date)"

# ------------------------------------------------------------------
echo ""
echo "=== STAGE 2: INSTRUCTION TUNING (~1.5 hrs) ==="
echo "Started: $(date)"
torchrun --nproc_per_node=8 scripts/train_stage2_ddp.py \
    --config $CONFIG \
    --stage1_ckpt checkpoints/stage1_final.pt
echo "Completed: $(date)"

# ------------------------------------------------------------------
echo ""
echo "=== STAGE 3: DOMAIN FINE-TUNE (~30 min) ==="
echo "Started: $(date)"
torchrun --nproc_per_node=8 scripts/train_stage3_ddp.py \
    --config $CONFIG \
    --stage2_ckpt checkpoints/stage2_final.pt
echo "Completed: $(date)"

# ------------------------------------------------------------------
echo ""
echo "=== PUSH CHECKPOINTS ==="
python3 scripts/push_checkpoints.py || echo "[WARN] Push failed, manual push needed"

# ------------------------------------------------------------------
echo ""
echo "=================================================================="
echo "TRAINING COMPLETE"
echo "Finished: $(date)"
echo "=================================================================="
echo ""
echo "Checkpoints:"
ls -lh checkpoints/stage*_final.pt 2>/dev/null || echo "  (none found)"
echo ""
echo "Next steps:"
echo "  1. Run benchmarks:  python3 scripts/run_benchmarks.py"
echo "  2. Export for ONNX:  python3 scripts/export_onnx.py"
echo "  3. Deploy:           python3 scripts/deploy.py"