#!/bin/bash
# Phase 2: v18s (mmBERT-base, continue v17s, 1 epoch) and v18L (ModernBERT-large from typed-decisions, 2 epochs), same data:
#   everything v17s saw + webgym new kinds (gym/clean3_*) + webgym DAgger (dagger/d_*). Then suites A, B, C and webgym held-out.
# Resumable: train.py resumes from <ckpt>/resume.pt; SKIP_BUILD=1 reuses the item files; STAGE=L|eval skips earlier stages.
cd /home/ckl/projects/S/laya && source env.sh
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True CKPT=1 COMPILE=0 LAYA_FMT=v3 LAYA_HEAD=768 MAX_TARGETS=40 FINAL_P=0.2
P=.venv/bin/python; O=finetune/out; I="bash finetune/infra.sh"; export INFRA_DIR=${INFRA_DIR:-/tmp/laya-infra}
SNAP=$(ls -d ~/.cache/huggingface/hub/models--convaiinnovations--laya/snapshots/*/typed-decisions | head -1)
bash finetune/infra.sh stop_s1; sleep 3   # chromium is shared with the other sessions: leave it running
cat $O/rollout2_cases.*.jsonl > $O/rollout2_cases.jsonl
cat $O/gym/*.jsonl > $O/gym_cases.jsonl
cat $O/dagger/d_*.jsonl > $O/dagger_gym_cases.jsonl 2>/dev/null || : > $O/dagger_gym_cases.jsonl
echo "== data: gym $(wc -l < $O/gym_cases.jsonl)  webgym-dagger $(wc -l < $O/dagger_gym_cases.jsonl)  m2w_sub $(wc -l < $O/m2w_sub_cases.jsonl)"
SRC="$O/done_cases.jsonl $O/step2_cases.jsonl $O/m2w_cases.jsonl $O/rollout_cases.jsonl $O/rollout2_cases.jsonl $O/nnetnav_cases.jsonl $O/gym_cases.jsonl $O/m2w_sub_cases.jsonl $O/dagger_gym_cases.jsonl"
train() {  # $1 ckpt name, $2 base dir, $3 items dir, $4 epochs
  mkdir -p $3
  if [ -z "$SKIP_BUILD" ] || [ ! -f $3/train_items.pt ]; then
    echo "== build $1"; LAYA_BASE=$2 $P finetune/build_items.py $O/pages.jsonl $O/cases.jsonl $3/ $SRC 2>&1 | grep -v Warn | tail -3
  fi
  echo "== train $1 ($4 epochs from $(basename $2))"; LAYA_BASE=$2 $P finetune/train.py $3/train_items.pt $O/$1 $4 2>&1 | grep --line-buffered -E "=== epoch|saved|Error|Traceback|resumed"
  [ -f $O/$1/model.safetensors ] || { echo "training of $1 interrupted; rerun (SKIP_BUILD=1) to resume"; exit 0; }
  echo "== calibrate + eval $1"; $P finetune/calibrate.py $O/pages.jsonl $3/eval_cases.jsonl $PWD/$O/$1 2>&1 | grep -vE "TileLang|Warn|warn|Fetch" | tail -1
  $P finetune/eval.py $O/pages.jsonl $3/eval_cases.jsonl $PWD/$O/$1 2>&1 | grep -E "operation acc|  live |  mind2web |  webgym"
}
[ "$STAGE" = "L" ] || [ "$STAGE" = "eval" ] || train laya-browser-v18s $PWD/$O/laya-browser-v17s $O/items_s 1
[ "$STAGE" = "eval" ] || train laya-browser-v18L $SNAP $O/items_L 1
$I start_chrome >/dev/null; SGL_MEM=0.5 $I start_sglang Qwen/Qwen3-8B-AWQ >/dev/null; bash finetune/webgym/serve.sh restart >/dev/null
until curl -s -m 3 http://127.0.0.1:30000/health >/dev/null; do sleep 5; done
KINDS=$($P -c "import sys; sys.path.insert(0,'finetune/webgym'); import spec; print(','.join(spec.KINDS))")
for CK in laya-browser-v18s laya-browser-v18L; do
  $I start_s1 $PWD/$O/$CK 60 >/dev/null
  echo "== $CK suite A x3"; (cd ../jev-ultrafast && REPEATS=3 SUITE_OUT=$PWD/../laya/$O/suiteA_$CK.json timeout 3600 .venv/bin/python ../laya/apps/browser_suite.py 2>&1 | grep -E "^==|per task")
  echo "== $CK suite B x3"; (cd ../jev-ultrafast && REPEATS=3 SUITE_OUT=$PWD/../laya/$O/suiteB_$CK.json timeout 3600 .venv/bin/python ../laya/apps/browser_suite_b.py 2>&1 | grep -E "^==|per task")
  [ -f apps/browser_suite_c.py ] && { echo "== $CK suite C x2"; (cd ../jev-ultrafast && REPEATS=2 SUITE_OUT=$PWD/../laya/$O/suiteC_$CK.json timeout 5400 .venv/bin/python ../laya/apps/browser_suite_c.py 2>&1 | grep -E "^==|per task|held-out"); }
  echo "== $CK webgym held-out ($KINDS)"; (cd ../jev-ultrafast && GYM_OUT=$PWD/../laya/$O/gym_eval_$CK.json timeout 5400 .venv/bin/python ../laya/finetune/webgym/eval_gym.py 10 $KINDS 60 2>&1 | grep -E "^==")
done
echo PHASE2_DONE