#!/usr/bin/env bash
# Crash-resilient full-scale training supervisor (stage1 → stage2).
#
# Both tsmm.train.stage1 and stage2 now support --resume: each saved checkpoint
# carries encoder/projector (+LoRA for stage2) + optimizer state + EMA loss, so
# a kill mid-run keeps all progress and the next --resume continues seamlessly.
#
# This supervisor loops "run --resume -> check step" until max_steps is reached.
# Each surviving iteration advances the step counter by >=1 (ckpt_every), so the
# run converges across multiple kills.
#
# Usage: scripts/train_full_supervisor.sh [stage1_max_steps] [stage2_max_steps]
set -uo pipefail   # no -e: a killed train sub-call must NOT abort the loop
cd "$(dirname "$0")/.."
export PYTHONPATH=src PYTHONUNBUFFERED=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

S1_MAX="${1:-31250}"   # 50万 align, bs=4 ga=8 => eff batch 32; 1 ep ≈ 15625, 2 ep = 31250
S2_MAX="${2:-9375}"    # 10万 sft,   bs=4 ga=8 => 1 ep ≈ 3125,  3 ep = 9375

run_until() {  # name  max  ckpt_dir  module  extra_args...
  local name=$1 max=$2 ckpt_dir=$3 module=$4; shift 4
  local extra=("$@")
  while :; do
    # Ask the training module's own _latest_step logic for the true highest-step
    # ckpt (it reads the 'step' field inside both step*.pt and final.pt).
    local cur latest_path
    latest_path=$(.venv/bin/python -c "
import sys; sys.path.insert(0,'src')
from tsmm.train.${module##*.} import _latest_step
import os
p=_latest_step('$ckpt_dir', '${name}_')
print(p if p else '')
" 2>/dev/null)
    if [ -n "$latest_path" ] && [ -f "$latest_path" ]; then
      cur=$(.venv/bin/python -c "
import torch
st=torch.load('$latest_path', map_location='cpu', weights_only=False)
print(int(st.get('step',0)))
" 2>/dev/null)
      cur=${cur:-0}
    else
      cur=0
    fi
    if [ "$cur" -ge "$max" ]; then
      echo "[sup] $name done: step $cur/$max (from $(basename "$latest_path"))"; return 0
    fi
    echo "[sup] $name at step $cur/$max — launching --resume $(date '+%H:%M:%S')"
    .venv/bin/python -m "$module" \
      --ckpt_dir "$ckpt_dir" --max_steps "$max" --resume "${extra[@]}"
    # loop: re-check (train may have been killed or completed)
  done
}

echo "[sup] === TRAIN START $(date '+%H:%M:%S') ==="
echo "[sup] stage1 target=$S1_MAX  stage2 target=$S2_MAX"

run_until stage1 "$S1_MAX" checkpoints/stage1 tsmm.train.stage1 \
  --data data/align.jsonl --bs 4 --grad_accum 8 --ctx 512 --lr 1e-4 \
  --lambda_contrast 0.1 --perturb_std 0.1 --log_every 50 --ckpt_every 2000 --tensorboard

run_until stage2 "$S2_MAX" checkpoints/stage2 tsmm.train.stage2 \
  --data data/sft.jsonl --stage1_ckpt checkpoints/stage1/stage1_final.pt \
  --bs 4 --grad_accum 8 --ctx 1024 --lr 1e-4 --lora_r 16 \
  --log_every 50 --ckpt_every 2000 --tensorboard

echo "[sup] === TRAIN ALL DONE $(date '+%H:%M:%S') ==="
ls -la checkpoints/stage1/stage1_final.pt checkpoints/stage2/stage2_final.pt 2>/dev/null