#!/usr/bin/env bash
# rf.py as written: rectified-flow DiT, 100 epochs on MNIST (the --cifar setting needs more than 16 GB of GPU memory
# at its fixed batch size of 256 and is not run). wandb disabled; the
# per-epoch loss by time bin printed by rf.py is the record. Datasets download into ./data on first use.
ROW=/var/tmp/kv-repro/046-cloneofsimo-minRF
export HOME=$ROW/data/home WANDB_MODE=disabled
for ds in mnist; do
  echo "KVERITAS_PHASE name=$ds"
  flag=""; [ $ds = cifar ] && flag=--cifar
  python rf.py $flag 2>&1 | tr '\r' '\n' | grep -v "it/s\]\|s/it\]\|^[0-9]*it \[" | tee $ROW/data/$ds.log
  python - $ROW/data/$ds.log $ds <<'PY'
import collections, re, sys
ds = sys.argv[2]
ep = collections.defaultdict(dict)
for e, b, v in re.findall(r"Epoch: (\d+), (\d) range loss: ([0-9.eE+-]+)", open(sys.argv[1]).read()):
    ep[int(e)][int(b)] = float(v)
if len(ep) != 100:
    sys.exit(f"{ds}: expected 100 epochs, got {len(ep)}")
for e in sorted(ep):
    print(f"KVERITAS_METRIC name={ds}_mean_bin_loss value={sum(ep[e].values()) / 10:.6f} step={e}")
last = ep[99]
print(f"KVERITAS_CLAIM metric={ds}_final_mean_bin_loss value={sum(last.values()) / 10:.6f}")
for b in (0, 5, 9):
    print(f"KVERITAS_CLAIM metric={ds}_final_bin{b}_loss value={last[b]:.6f}")
PY
  [ $? -eq 0 ] || exit 1
done
