#!/bin/bash
# Judge-in-loop ABC (natural Llama-8B base). Seed once via judge@50->100, train A/B/C
# with REWARD_MODE x BRIER_FROM_STATED=1 (martingale=judge, Brier=stated-P).
set -uo pipefail
cd /workspace/Eval-Reasoning-Consistency
source /workspace/.orkey
VENV=/workspace/martingale_training_env/bin/python
TM=/workspace/Eval-Reasoning-Consistency/training/trained_models
BASE=/tmp/models_tmp/Llama-3.1-8B-Instruct
log(){ echo "[$(date -u +%H:%M:%S)] $*"; }

# 1. serve base, generate the shared judge-seed over N=2000 (needs sglang)
log "serving base for seed-gen"
curl -s -m3 http://127.0.0.1:30000/v1/models >/dev/null 2>&1 || bash /tmp/start_sgLlama.sh
for i in $(seq 1 60); do curl -s -m3 http://127.0.0.1:30000/v1/models >/dev/null 2>&1 && break; sleep 10; done
if [ ! -f /tmp/seed_judge_base.json ]; then
  log "gen_seed_judge over llama_train_q2000.json ..."
  $VENV /tmp/gen_seed_judge.py /workspace/llama_train_q2000.json BASEJUDGE /tmp/seed_judge_base.json
fi
N=$($VENV -c 'import json;print(len(json.load(open("/tmp/seed_judge_base.json"))))')
log "SEED_READY n=$N"

# 2. free the GPU (trainer needs it; sglang reserves 85%)
log "killing sglang to free GPU for training"
pkill -9 -f sglang 2>/dev/null; sleep 8

# 3. train A/B/C on the shared judge-seed
for ARM in A B C; do
  RID=Llama8B-JUDGE-$ARM
  RUNDIR=data/runs/batch-martingale-training/run-$RID
  mkdir -p "$RUNDIR"; cp /tmp/seed_judge_base.json "$RUNDIR/sft-data.json"
  rm -f "$TM/training_state.json"; rm -rf "$TM/epoch_0"
  export PATH=/workspace/martingale_training_env/bin:$PATH
  export PY_EXEC=$VENV DIR_NAME=batch-martingale-training EXCLUDE_BASE_EVAL=1
  export RUN_ID=$RID MODEL_BASE_DIR=$BASE
  export KL_SAFE_REWARD=1 KL_BETA=0.1 INFO_TERM=1 INFO_COEF=1.0 INFO_EPS=0.02
  export TRAIN_LR=3e-5 MAX_EPOCH=1 TRAIN_CONTEXT_LEN=4092 LORA=1 PPO_CLIP_EPS=0.0
  export REWARD_MODE=$ARM LAMBDA_C=0.5 BRIER_FROM_STATED=1
  log "### TRAIN $ARM ($RID) REWARD_MODE=$ARM BRIER_FROM_STATED=1 ###"
  $VENV -u -m training.sft_product_based > /tmp/train_judge_$ARM.log 2>&1
  if [ -f "$TM/epoch_0/adapter_model.safetensors" ]; then
    rm -rf /data/jobs/abcjudge_$ARM; cp -r "$TM/epoch_0" /data/jobs/abcjudge_$ARM
    log "### $ARM DONE -> /data/jobs/abcjudge_$ARM ###"
  else
    log "### $ARM FAILED (no adapter); tail log:"; tail -5 /tmp/train_judge_$ARM.log
  fi
done
log "JUDGE_ABC_DONE"