#!/usr/bin/env bash
# Training stages on the pod (classifier-plan §3.2). Run each inside tmux; every training stage resumes from last/.
# The local ledger (`python -m training.ledger check`) must say "go" before each stage starts (runbook.sh does this).
#   run.sh setup            venv, pinned installs, base downloads, check_env, selftest (bootstrap.sh)
#   run.sh smoke            100 steps on 1k posts; stop at 50 and resume; forced step on the 8 longest; predict repeat
#   run.sh screen           4 arms (B-last, I-last, B-mean, I-mean), 1 epoch (seeded 8k subset if train > 8k), choose
#   run.sh final ARM        all train, ${EPOCHS:-2} epochs, val every 200 steps, best/ = lowest mean val log-loss
#   run.sh evaluate RUN     val (reload) predictions, temperatures, test (repeat 2), robustness predictions
#   run.sh cleanup          delete /root/data (post text) and write /workspace/CLEAN
# Data lives on the container disk only (/root/data/bundle, shipped by sync.sh up); runs on /workspace/runs.
set -euo pipefail
W=/workspace
cd "$W/code"
DATA=${DATA:-/root/data/bundle}
M=${MODELS:-/root/models}
RUNS=$W/runs
export HF_HOME=/root/hf
activate() { . "${VENV:-/root/venv}/bin/activate"; }

arm_model() { case "$1" in B-*) echo "$M/gemma-4-12B" ;; I-*) echo "$M/gemma-4-12B-it" ;; *) echo "bad arm $1" >&2; exit 2 ;; esac; }
arm_pool() { case "$1" in *-last) echo last ;; *-mean) echo mean ;; *) echo "bad arm $1" >&2; exit 2 ;; esac; }
arm_checkpoint() { case "$1" in B-*) echo base ;; I-*) echo it ;; esac; }

# train ARM OUT HOURS [extra args]: one resumable training process under timeout (SIGINT -> saved resume point).
train() { local arm=$1 out=$2 hours=$3; shift 3
  timeout --signal=INT -k 300 "$(python3 -c "print(int(($hours + 0.25) * 3600))")" \
    python -m training.train --data "$DATA" --model "$(arm_model "$arm")" --out "$out" --arm "$arm" \
      --pool "$(arm_pool "$arm")" --checkpoint "$(arm_checkpoint "$arm")" --max-hours "$hours" --device cuda "$@"; }

predict() { timeout -k 60 "${PREDICT_TIMEOUT:-2700}" python -m training.predict --device cuda "$@"; }

case "${1:-}" in
  setup)
    bash training/runpod/bootstrap.sh ;;
  smoke)
    activate
    out=$RUNS/smoke
    rc=0
    train B-last "$out" 0.6 --subset 1000 --epochs 4 --max-steps 100 --eval-every 50 --val-limit 300 --log-every 5 \
      --force-longest 8 --stop-after 50 || rc=$?
    [ $rc = 3 ] || { echo "smoke: expected exit 3 at the planned stop, got $rc"; exit 1; }
    train B-last "$out" 0.6 --subset 1000 --epochs 4 --max-steps 100 --eval-every 50 --val-limit 300 --log-every 5 \
      --resume
    predict --run "$out" --split val --repeat 2 --out "$out/predictions/val-repeat.jsonl" | tee "$out/predict_repeat.json"
    grep '"event": *"train"' "$out/train_log.jsonl" | tail -1 ;;
  screen)
    activate
    n=$(wc -l < "$DATA/train.jsonl")
    sub=0; [ "$n" -gt 8000 ] && sub=8000
    runs=()
    for arm in B-last I-last B-mean I-mean; do
      out=$RUNS/screen-$arm
      train "$arm" "$out" "${SCREEN_HOURS:-0.4}" --subset "$sub" --epochs 1 --resume
      runs+=("$out")
    done
    python -m training.screen "${runs[@]}" --data "$DATA" --out "$RUNS/screen.json" ;;
  final)
    activate
    arm=${2:?arm: B-last, I-last, B-mean or I-mean}
    train "$arm" "$RUNS/final-$arm" "${FINAL_HOURS:-1.5}" --epochs "${EPOCHS:-2}" --eval-every 200 --resume ;;
  evaluate)
    activate
    run=${2:?run directory}; p=$run/predictions
    mkdir -p "$p"
    # Temperatures are fitted on predictions from the saved best/ checkpoint (what node12 will serve).
    predict --run "$run" --split val --out "$p/val-reload.jsonl" | tee "$run/val_predict.json"
    [ -f "$run/temperatures.json" ] || python -m training.calibrate fit "$p/val-reload.jsonl" --data "$DATA" \
      --out "$run/temperatures.json"
    predict --run "$run" --split test --repeat 2 --out "$p/test.jsonl" | tee "$run/test_predict.json"
    rb=$DATA/robustness
    [ -f "$rb/robustness.jsonl" ] || python -m training.robustness build --data "$DATA" --out "$rb"
    predict --run "$run" --input "$rb/robustness.jsonl" --out "$p/robustness.jsonl" | tee "$run/robustness_predict.json" ;;
  cleanup)
    rm -rf /root/data
    [ ! -e /root/data ] || { echo "cleanup: /root/data still exists"; exit 1; }
    date -u +%Y-%m-%dT%H:%M:%SZ > "$W/CLEAN"
    echo "CLEAN written" ;;
  *) sed -n 2,10p "$0"; exit 2 ;;
esac