#!/bin/bash
# LR check (S, 20 min) -> Engram A/B (L, 60 min, equal wall-clock) -> 23h main run with the winners.
# Runs in the background; every decision is printed to $TA_DATA/logs/pipeline.log.
set -euo pipefail
cd "$(dirname "$0")/.."
source env.sh
R=$TA_DATA/runs
mkdir -p $R
log() { echo "[$(date '+%F %T')] $*" | tee -a $TA_DATA/logs/pipeline.log; }

common="--tokens 1e10 --warmup 50 --eval_every 50 --ckpt_minutes 999 --snapshot_tokens 1e12"

log "waiting for synth_grounded (part of the stable mixture)"
until ls $TA_DATA/tok/train/synth_grounded/*.bin >/dev/null 2>&1; do sleep 30; done
log "phase 1: LR check at size S (3 concurrent, 20 min)"
for lr in 1.5e-3 3e-3 6e-3; do
  [ -f $R/lr_$lr/ckpt.pt ] || $TA_PY scripts/train.py --size S --engram 0 --lr $lr --batch_tokens 131072 \
      --stop_minutes 20 $common --out $R/lr_$lr > $R/lr_$lr.out 2>&1 &
done
wait
$TA_PY scripts/compare_runs.py $R/lr_1.5e-3 $R/lr_3e-3 $R/lr_6e-3 | tee -a $TA_DATA/logs/pipeline.log
LR=$($TA_PY - <<EOF
import json
best=None
for lr in ["1.5e-3","3e-3","6e-3"]:
    recs=[json.loads(l) for l in open("$R/lr_"+lr+"/log.jsonl") if '"val"' in l]
    v=recs[-1]["val"]; m=sum(v.values())/len(v)
    if best is None or m<best[0]: best=(m,lr)
print(best[1])
EOF
)
log "best short-run lr=$LR"
if [ -n "${LR_OVERRIDE:-}" ]; then LR=$LR_OVERRIDE; log "LR_OVERRIDE -> lr=$LR"; fi

log "phase 2: Engram A/B at size L (2 concurrent, 60 min, equal wall-clock)"
for e in 0 1; do
  [ -f $R/ab_engram$e/ckpt.pt ] || $TA_PY scripts/train.py --size L --engram $e --lr $LR --warmup 100 \
      --stop_minutes 60 --tokens 1e10 --eval_every 50 --ckpt_minutes 999 --snapshot_tokens 1e12 \
      --out $R/ab_engram$e > $R/ab_engram$e.out 2>&1 &
done
wait
$TA_PY scripts/compare_runs.py $R/ab_engram0 $R/ab_engram1 | tee -a $TA_DATA/logs/pipeline.log
ENGRAM=$($TA_PY - <<EOF
import json
def last(d):
    recs=[json.loads(l) for l in open(d+"/log.jsonl") if '"val"' in l]
    v=recs[-1]["val"]; return sum(v.values())/len(v)
# equal wall-clock: whichever reached the lower val loss in the same 60 minutes wins
print(1 if last("$R/ab_engram1") < last("$R/ab_engram0") else 0)
EOF
)
log "picked engram=$ENGRAM"

log "waiting for synthetic sources"
n_parts() { ls $TA_DATA/tok/train/$1/*.bin 2>/dev/null | wc -l; }
until [ "$(n_parts synth_grounded)" -ge 8 ] && [ "$(n_parts synth_reasoning)" -ge 8 ]; do sleep 30; done
log "phase 3: main run, size L, engram=$ENGRAM, lr=$LR, 23h budget with automatic decay, decay at T=4096"
$TA_PY scripts/train.py --size L --engram $ENGRAM --lr $LR --tokens 1e10 --max_minutes 1380 --decay_T 4096 \
    --eval_every 250 --snapshot_tokens 5e8 --out $R/main > $R/main.out 2>&1
log "main run finished"