#!/usr/bin/env bash
# Training stages on the pod (spec section 6). Run each inside tmux; every stage resumes from last/.
# The local ledger (classifier.ledger check) must say "go" before each stage starts.
#   run.sh smoke          stage 1: 100 steps on 1k posts, stop at 50 and resume, forced step on the 8 longest
#   KERNELS=torch run.sh ... runs a stage on the PyTorch path (the spec's pre-declared fallback)
#   run.sh curve          stage 2: seeded 8k train posts, up to 3 epochs, dev every 50 steps
#   run.sh full           stage 3: all train posts, 2 epochs, dev every 200 steps
#   run.sh evaluate RUN   stage 4: dev, benchmark (repeated) and robustness predictions; dev temperatures
#   run.sh pool RUN       stage 5: forward-only scores for the local-context pool
#   run.sh e1 | e2        optional extras: post-trained checkpoint on 8k | 2k curve point
set -euo pipefail
W=/workspace
cd "$W/code"
. "${VENV:-/root/venv}/bin/activate"
export HF_HOME=/root/hf
DATA=$W/data/classifier-data-20261017
M=${MODELS:-/root/models}             # container disk (bootstrap.sh)
BASE=$M/Qwen3.5-9B-Base
train() { local out=$1 hours=$2; shift 2
  timeout --signal=INT -k 300 "$(python3 -c "print(int(($hours + 0.25) * 3600))")" \
    python -m classifier.train --data "$DATA" --model "${MODEL:-$BASE}" --out "$out" --max-hours "$hours" \
      --kernels "${KERNELS:-fast}" "$@"; }

case "${1:-}" in
  smoke)
    out=$W/runs/smoke
    rc=0; train "$out" 0.6 --subset 1000 --epochs 4 --max-steps 100 --eval-every 50 --dev-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 "$out" 0.6 --subset 1000 --epochs 4 --max-steps 100 --eval-every 50 --dev-limit 300 --log-every 5 --resume
    timeout -k 60 900 python -m classifier.predict --kernels "${KERNELS:-fast}" --run "$out" --split dev --limit 200 --repeat 2 \
      --out "$out/predictions/dev-repeat.jsonl" | tee "$out/predict_repeat.json" ;;
  curve)
    train "$W/runs/curve-8k" 1.0 --subset 8000 --epochs 3 --eval-every 50 --resume ;;
  full)
    train "$W/runs/full" 2.5 --epochs 2 --eval-every 200 --resume ;;
  evaluate)
    run=${2:?run directory}; p=$run/predictions
    # Temperatures come from the training-time predictions of the best checkpoint (predictions/dev.jsonl).
    # dev-reload.jsonl re-predicts dev from the saved checkpoint: classifier_eval determinism compares the two.
    timeout -k 60 2700 python -m classifier.predict --kernels "${KERNELS:-fast}" --run "$run" --split dev --out "$p/dev-reload.jsonl" | tee "$run/dev_predict.json"
    python -m classifier.calibrate fit "$p/dev.jsonl" --data "$DATA" --out "$run/temperatures.json"
    timeout -k 60 2700 python -m classifier.predict --kernels "${KERNELS:-fast}" --run "$run" --split benchmark --repeat 2 --out "$p/benchmark.jsonl" \
      | tee "$run/benchmark_predict.json"
    timeout -k 60 1800 python -m classifier.predict --kernels "${KERNELS:-fast}" --run "$run" --input "$W/inputs/robustness.jsonl" --out "$p/robustness.jsonl" \
      | tee "$run/robustness_predict.json" ;;
  pool)
    run=${2:?run directory}
    for f in "$DATA"/pool/*.jsonl; do
      timeout -k 60 1800 python -m classifier.predict --kernels "${KERNELS:-fast}" --run "$run" --input "$f" --out "$run/predictions/pool-$(basename "$f")" \
        | tee -a "$run/pool_predict.json"
    done ;;
  e1)
    MODEL=$M/Qwen3.5-9B train "$W/runs/e1-instruct-8k" 1.0 --subset 8000 --epochs 3 --eval-every 50 --resume ;;
  e2)
    train "$W/runs/curve-2k" 0.3 --subset 2000 --epochs 3 --eval-every 25 --resume ;;
  *) sed -n 2,10p "$0"; exit 2 ;;
esac