#!/usr/bin/env bash
# P2: one full expert-iteration round, as an explicit sequence of RESUMABLE steps.
#
#   scripts/p2/run_round.sh <step> [<step>...]
#   scripts/p2/run_round.sh all            # in order; stops at the first failure
#   scripts/p2/run_round.sh print          # show every command, run nothing
#
# Steps (each is one process that takes and releases the repo heavy-job lock itself; NEVER run
# two at once, and never wrap the whole script in a lock: the children would deadlock on it):
#
#   problems     eval sets + decontaminated training pool + elaboration filter
#   base_eval    baseline: generate (cap 16384, n=8) + verify on miniF2F and MOBench
#   pool_gen     round-0 generation on the training pool
#   pool_verify  verify round 0 (with --feedback: correction prompts need the Lean errors)
#   corr_gen     correction round 1 from the failed round-0 samples
#   corr_verify  verify round 1
#   sft_data     build the SFT shard set (exact sampled tokens, completion-only loss)
#   train        P2-B's trainer (command supplied via $P2_TRAIN_CMD; not defined here)
#   export       P2-B's HF export of the trained weights (command via $P2_EXPORT_CMD)
#   post_eval    the same evaluation as base_eval on the trained weights
#   report       eval_report.py: pass@k, CIs, paired McNemar, correction lift
#
# Required env: RUN (run name); MAX_SEQ_LEN for sft_data. Optional: CAP (16384), N_EVAL (8), N_POOL (8), WORKERS (8),
# TP (4), CAP_REF (40960, set to also run the reference-cap baseline).
# Host rules (/apps/CLAUDE.md) are enforced by the python entry points (preflight + flock +
# <= 8 Lean workers + scratch under /apps); this script only sequences them.
set -euo pipefail

REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
cd "$REPO"
: "${RUN:?set RUN=<run name>}"
CAP="${CAP:-16384}"
N_EVAL="${N_EVAL:-8}"
N_POOL="${N_POOL:-8}"
WORKERS="${WORKERS:-8}"
TP="${TP:-4}"
D="$REPO/.tmp/p2/$RUN"
PY="$REPO/.venv/bin/python"
PYR="$REPO/.venv-rollout/bin/python"
export TMPDIR="$D/scratch" HF_HOME="${HF_HOME:-/apps/development/rlhf/hf_cache}" HF_HUB_DISABLE_XET=1
MODEL_BASE="${MODEL_BASE:-/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/dfd02e6271a58375dfbf3ece0175277cf6b6a89a}"
MODEL_SFT="${MODEL_SFT:-$D/export}"
PRINT_ONLY=0

run() {  # print the command; execute unless printing
  echo "+ $*" >&2
  [ "$PRINT_ONLY" = 1 ] || "$@"
}

gen()    { # gen <model> <problems> <out> <n> <cap> [extra...]
  local model="$1" probs="$2" out="$3" n="$4" cap="$5"; shift 5
  run "$PYR" extras/vllm_rollout/lean_gen.py --model "$model" --problems "$probs" --out "$out" --n "$n" \
      --cap-tokens "$cap" --round 0 --tp "$TP" --seed 0 --wait-lock "$@"
}
verify() { # verify <gens> <problems> <out> [extra...]
  local g="$1" p="$2" o="$3"; shift 3
  run "$PY" scripts/p2/verify_generations.py --generations "$g" --problems "$p" --out "$o" --workers "$WORKERS" "$@"
}

step_problems() {
  mkdir -p "$D/problems" "$TMPDIR"
  run "$PY"  scripts/p2/build_problems.py --phase eval        --out-dir "$D/problems"
  run "$PYR" scripts/p2/build_problems.py --phase pool-sample --out-dir "$D/problems" --n 2000 --candidates 5000
  run "$PY"  scripts/p2/build_problems.py --phase check       --out-dir "$D/problems" --n 2000 --workers "$WORKERS"
}
export_rev() {  # content id of an exported model (INTEGRATOR FIX: lean_gen refuses a non-snapshot dir without --model-rev)
  [ "$PRINT_ONLY" = 1 ] && { echo "<export-rev>"; return; }
  echo "sft-$(cd "$1" && cat $(ls model*.safetensors | sort) | sha256sum | cut -c1-16)"
}
eval_one() { # eval_one <tag> <model> <cap>
  local tag="$1" model="$2" cap="$3" rev=()
  [ "$model" = "$MODEL_BASE" ] || rev=(--model-rev "$(export_rev "$model")")
  for set in minif2f mobench; do
    gen "$model" "$D/problems/$set.jsonl" "$D/$tag/gen_${set}_c${cap}.jsonl" "$N_EVAL" "$cap" "${rev[@]}"
    verify "$D/$tag/gen_${set}_c${cap}.jsonl" "$D/problems/$set.jsonl" "$D/$tag/ver_${set}_c${cap}.jsonl"
  done
}
step_base_eval() {
  eval_one base "$MODEL_BASE" "$CAP"
  [ -z "${CAP_REF:-}" ] || eval_one base "$MODEL_BASE" "$CAP_REF"
}
step_pool_gen()    { gen "$MODEL_BASE" "$D/problems/train_pool.jsonl" "$D/pool/gen_r0.jsonl" "$N_POOL" "$CAP"; }
step_pool_verify() { verify "$D/pool/gen_r0.jsonl" "$D/problems/train_pool.jsonl" "$D/pool/ver_r0.jsonl" --feedback; }
step_corr_gen() {
  run "$PYR" extras/vllm_rollout/lean_gen.py --model "$MODEL_BASE" --problems "$D/problems/train_pool.jsonl" \
      --out "$D/pool/gen_r1.jsonl" --n 1 --cap-tokens "$CAP" --round 1 --tp "$TP" --seed 0 --wait-lock \
      --corrections-from "$D/pool/ver_r0.jsonl" "$D/pool/gen_r0.jsonl"
}
step_corr_verify() { verify "$D/pool/gen_r1.jsonl" "$D/problems/train_pool.jsonl" "$D/pool/ver_r1.jsonl"; }
step_sft_data() {
  # No default on purpose: the longest row the trainer can hold (prompt + completion; correction
  # rows carry the failed attempt in the prompt) is a memory fact P2-B measures, not a constant.
  [ "$PRINT_ONLY" = 1 ] || : "${MAX_SEQ_LEN:?set MAX_SEQ_LEN to the longest sequence the trainer fits (P2-B probe)}"
  run "$PY" scripts/p2/build_sft_data.py \
      --generations "$D/pool/gen_r0.jsonl" "$D/pool/gen_r1.jsonl" \
      --verifications "$D/pool/ver_r0.jsonl" "$D/pool/ver_r1.jsonl" \
      --problems "$D/problems/train_pool.jsonl" \
      --eval-problems "$D/problems/minif2f.jsonl" "$D/problems/mobench.jsonl" \
      --cap-tokens "$CAP" --out-dir "$D/sft_data" --max-seq-len "${MAX_SEQ_LEN:-<MAX_SEQ_LEN>}"
}
step_train() {
  [ "$PRINT_ONLY" = 1 ] || : "${P2_TRAIN_CMD:?P2-B supplies the training command (reads $D/sft_data/shards)}"
  run bash -c "${P2_TRAIN_CMD:-<P2_TRAIN_CMD>}"
}
step_export() {
  [ "$PRINT_ONLY" = 1 ] || : "${P2_EXPORT_CMD:?P2-B supplies the HF export command (writes $MODEL_SFT)}"
  run bash -c "${P2_EXPORT_CMD:-<P2_EXPORT_CMD>}"
}
step_post_eval() {
  eval_one sft "$MODEL_SFT" "$CAP"
  [ -z "${CAP_REF:-}" ] || eval_one sft "$MODEL_SFT" "$CAP_REF"
}
step_report() {
  local runs=() cap
  for tag in base sft; do
    local g="" v=""
    for cap in "$CAP" ${CAP_REF:-}; do
      for set in minif2f mobench; do g+="${g:+,}$D/$tag/gen_${set}_c${cap}.jsonl"; v+="${v:+,}$D/$tag/ver_${set}_c${cap}.jsonl"; done
    done
    runs+=(--run "$tag:$g:$v")
  done
  run "$PY" scripts/p2/eval_report.py --problems "$D/problems/minif2f.jsonl" "$D/problems/mobench.jsonl" \
      "${runs[@]}" --baseline base --candidate sft --out-dir "$D/report"
}

[ "$#" -ge 1 ] || { sed -n '2,30p' "$0"; exit 2; }
if [ "$1" = print ]; then PRINT_ONLY=1; set -- all; fi
steps=("$@")
[ "${steps[0]}" != all ] || steps=(problems base_eval pool_gen pool_verify corr_gen corr_verify sft_data train export post_eval report)
for s in "${steps[@]}"; do
  echo "=== step: $s ($(date '+%F %T'))" >&2
  "step_$s"
done