#!/usr/bin/env bash
# P2 integrator: ONE tiny end-to-end loop on the real model, serialized under the heavy lock.
#
#   scripts/p2/smoke_loop.sh <step> [<step>...]      steps in order: see STEPS below
#   scripts/p2/smoke_loop.sh print                   show every command, run nothing
#
# Purpose: prove the PLUMBING (keys, tokens, files, the export) with numbers that mean nothing:
# 16 short training statements, 8 miniF2F problems, cap 4096. It is NOT the experiment.
#
# Host rules (/apps/CLAUDE.md): every heavy step is preceded by `preflight` (>=45 GB MemAvailable,
# load < 2x cores, disk queue < 64; strict, raises) and by `waitfree` (poll the repo heavy-job flock
# until nobody holds it). The python entry points then take the lock themselves for the whole
# step; this script never holds it (children would deadlock on it). Scratch is under
# .tmp/p2/p2i, TMPDIR too; nothing goes to / or /tmp.
#
# STEPS: problems gen_r0 ver_r0 gen_r1 ver_r1 sft_data seams train_a train_resume export gen_post gen_base ver_eval report cleanup
set -euo pipefail

REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
cd "$REPO"
RUN="${RUN:-p2i}"
D="$REPO/.tmp/p2/$RUN"
PY="$REPO/.venv/bin/python"
PYR="$REPO/.venv-rollout/bin/python"
CAP="${CAP:-4096}"
TP="${TP:-4}"
WORKERS="${WORKERS:-8}"
MAX_SEQ_LEN="${MAX_SEQ_LEN:-9216}"          # correction rows carry the failed attempt: <= cap + prompt + errors
MODEL_BASE="${MODEL_BASE:-/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/dfd02e6271a58375dfbf3ece0175277cf6b6a89a}"
LOCK="$REPO/.align_heavy_job.lock"
export TMPDIR="$D/scratch" HF_HOME="${HF_HOME:-/apps/development/rlhf/hf_cache}" HF_HUB_DISABLE_XET=1
PRINT_ONLY=0
mkdir -p "$TMPDIR" "$D/problems" "$D/logs"

run() { echo "+ $*" >&2; [ "$PRINT_ONLY" = 1 ] || "$@"; }

preflight() {  # strict host preflight; prints the numbers it saw
  [ "$PRINT_ONLY" = 1 ] && { echo "+ preflight" >&2; return; }
  "$PY" -c "from align.utils.guard import preflight_host_check as p; import json; print('preflight', json.dumps(p(strict=True)))" >&2
}
waitfree() {   # block until the heavy-job lock is free (probe + release; the job takes it itself)
  [ "$PRINT_ONLY" = 1 ] && { echo "+ waitfree" >&2; return; }
  local t0=$SECONDS
  while ! flock -n "$LOCK" true; do
    echo "waiting for heavy lock held by: $(cat "$LOCK" 2>/dev/null)" >&2; sleep 60
  done
  echo "lock free after $((SECONDS - t0)) s" >&2
}
heavy() { waitfree; preflight; run "$@"; }
heavy_log() {  # heavy_log <logfile> cmd... : as `heavy`, output also appended to the log (pipefail keeps the exit code)
  local log="$1"; shift; waitfree; preflight; echo "+ $*" >&2
  [ "$PRINT_ONLY" = 1 ] || "$@" 2>&1 | tee -a "$log"
}

step_problems() {
  # Real eval set + a 16-problem slice of the decontaminated 2,000-problem pool (the 16 SHORTEST
  # statements: this is a plumbing test and wants some accepted proofs), 8 miniF2F problems.
  [ "$PRINT_ONLY" = 1 ] && { echo "+ (slice problems -> $D/problems)" >&2; return; }
  "$PY" - "$REPO/.tmp/p2/problems" "$D/problems" <<'PYEOF'
import json, sys
src, dst = sys.argv[1:3]
tr = [json.loads(l) for l in open(f"{src}/train_pool.jsonl")]
mf = [json.loads(l) for l in open(f"{src}/minif2f.jsonl")]
tr.sort(key=lambda r: (len(r["formal_statement"]), r["problem_id"]))
w = lambda p, rows: open(p, "w").write("".join(json.dumps(r, ensure_ascii=False) + "\n" for r in rows))
w(f"{dst}/train16.jsonl", tr[:16]); w(f"{dst}/minif2f8.jsonl", mf[:8]); w(f"{dst}/minif2f.jsonl", mf)
w(f"{dst}/mobench.jsonl", [json.loads(l) for l in open(f"{src}/mobench.jsonl")])
PYEOF
}
gen() { # gen <model> <problems> <out> <n> <cap> [extra...]
  local model="$1" probs="$2" out="$3" n="$4" cap="$5"; shift 5
  heavy "$PYR" extras/vllm_rollout/lean_gen.py --model "$model" --problems "$probs" --out "$out" --n "$n" \
      --cap-tokens "$cap" --tp "$TP" --seed 0 --wait-lock "$@"
}
verify() { # verify <gens> <problems> <out> [extra...]
  local g="$1" p="$2" o="$3"; shift 3
  heavy "$PY" scripts/p2/verify_generations.py --generations "$g" --problems "$p" --out "$o" --workers "$WORKERS" "$@"
}
step_gen_r0()  { gen "$MODEL_BASE" "$D/problems/train16.jsonl" "$D/pool/gen_r0.jsonl" 4 "$CAP" --round 0; }
step_ver_r0()  { verify "$D/pool/gen_r0.jsonl" "$D/problems/train16.jsonl" "$D/pool/ver_r0.jsonl" --feedback; }
step_gen_r1()  { gen "$MODEL_BASE" "$D/problems/train16.jsonl" "$D/pool/gen_r1.jsonl" "${N_CORR:-1}" "$CAP" --round 1 \
                   --corrections-from "$D/pool/ver_r0.jsonl" "$D/pool/gen_r0.jsonl"; }
step_ver_r1()  { verify "$D/pool/gen_r1.jsonl" "$D/problems/train16.jsonl" "$D/pool/ver_r1.jsonl"; }
step_sft_data() {
  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/train16.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" --held-frac 0 ${SFT_EXTRA:-}
}
step_seams() {
  run "$PY" scripts/p2/check_seams.py keys --generations "$D/pool/gen_r0.jsonl" --verifications "$D/pool/ver_r0.jsonl"
  [ ! -s "$D/pool/gen_r1.jsonl" ] || run "$PY" scripts/p2/check_seams.py keys --generations "$D/pool/gen_r1.jsonl" --verifications "$D/pool/ver_r1.jsonl"
  run "$PY" scripts/p2/check_seams.py sft --generations "$D/pool/gen_r0.jsonl" "$D/pool/gen_r1.jsonl" \
      --verifications "$D/pool/ver_r0.jsonl" "$D/pool/ver_r1.jsonl" --shards "$D/sft_data/shards"
}

# ---- training: 4 ranks, LoRA, configs/p2_sft.toml + overrides --------------------------------
TRAIN_SEQ_LEN="${TRAIN_SEQ_LEN:-9216}"; TRAIN_EPOCHS="${TRAIN_EPOCHS:-1}"; SAVE_EVERY="${SAVE_EVERY:-20}"
train_cmd() {
  echo "env PYTHONPATH=$REPO $PY -m torch.distributed.run --standalone --nproc_per_node=4 scripts/sft.py configs/p2_sft.toml" \
       "--data.shard_dir=$D/sft_data/shards --data.manifest_path=$D/sft_data/manifest.json" \
       "--paths.checkpoint_dir=$D/sft --train.seq_len=$TRAIN_SEQ_LEN --train.epochs=$TRAIN_EPOCHS" \
       "--train.save_every_micro_steps=$SAVE_EVERY"
}
step_train_a()      { heavy_log "$D/logs/train.log" bash -c "$(train_cmd)"; }   # SIGTERM'd by the operator, see docs/P2.md
step_train_resume() { heavy_log "$D/logs/train_resume.log" bash -c "$(train_cmd)"; }
step_export() {  # merge adapter + HF export + vLLM reload + first-token logprob seam (keeps the export)
  local ck="<latest micro_N>"
  if [ "$PRINT_ONLY" != 1 ]; then ck="$(ls -d "$D"/sft/*/checkpoints/micro_????????  2>/dev/null | sort -V | tail -1)"; : "${ck:?no checkpoint under $D/sft}"; fi
  heavy "$PY" scripts/qwen3_export_check.py run --ckpt "$ck" --out-dir "$D/export" --work "$D/export_check" --r 64 --alpha 128 --keep
}
export_rev() {  # content id of an exported model: lean_gen cannot infer a revision from a non-snapshot dir (it refuses)
  [ "$PRINT_ONLY" = 1 ] && { echo "<export-rev>"; return; }
  echo "sft-$(cd "$1" && cat $(ls model*.safetensors | sort) | sha256sum | cut -c1-16)"
}
step_gen_post() { gen "$D/export" "$D/problems/minif2f8.jsonl" "$D/sft/gen_minif2f8.jsonl" 2 "$CAP" --round 0 --model-rev "$(export_rev "$D/export")"; }
step_gen_base() { gen "$MODEL_BASE" "$D/problems/minif2f8.jsonl" "$D/base/gen_minif2f8.jsonl" 2 "$CAP" --round 0; }
step_ver_eval() {
  verify "$D/base/gen_minif2f8.jsonl" "$D/problems/minif2f8.jsonl" "$D/base/ver_minif2f8.jsonl"
  verify "$D/sft/gen_minif2f8.jsonl"  "$D/problems/minif2f8.jsonl" "$D/sft/ver_minif2f8.jsonl"
}
step_report() {
  run "$PY" scripts/p2/eval_report.py --problems "$D/problems/minif2f8.jsonl" \
      --run "base:$D/base/gen_minif2f8.jsonl:$D/base/ver_minif2f8.jsonl" \
      --run "sft:$D/sft/gen_minif2f8.jsonl:$D/sft/ver_minif2f8.jsonl" \
      --baseline base --candidate sft --out-dir "$D/report" --ks 1 2
}
step_cleanup() { run rm -rf "$D/export" "$D/scratch"; }

[ "$#" -ge 1 ] || { sed -n '2,10p' "$0"; exit 2; }
if [ "$1" = print ]; then PRINT_ONLY=1; set -- problems gen_r0 ver_r0 gen_r1 ver_r1 sft_data seams train_a train_resume export gen_post gen_base ver_eval report cleanup; fi
for s in "$@"; do echo "=== step: $s ($(date '+%F %T'))" >&2; "step_$s"; done