#!/bin/bash
# User order: baseline -> step1200 -> A/B scores -> then 200/400/800 -> table.
set -uo pipefail
exec > /opt/work/sweep3.log 2>&1
CKPTDIR=/opt/work/fish-speech/results/s2pro_egy_fast/checkpoints
PY=/opt/work/fish-speech/.venv/bin/python
MEGA=/opt/work/scripts/eval_mega.jsonl
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
cd /opt/work/fish-speech
source /opt/work/scripts/soniox.env
echo "=== chain_sweep3 $(date) ==="

# wait for any in-flight synth (baseline) to finish
while ps aux | grep synth_eval.py | grep -v grep >/dev/null; do sleep 10; done
echo "== mega: baseline (fill if incomplete) =="
$PY /opt/work/scripts/synth_eval.py --sentences $MEGA --outdir /opt/work/eval/mega_baseline --skip-existing

echo "== mega: v2 step 1200 =="
$PY /opt/work/scripts/synth_eval.py --sentences $MEGA --outdir /opt/work/eval/mega_step1200 \
    --lora-ckpt "$CKPTDIR/step_000001200.ckpt" --lora-config r_32_alpha_16_fast

/opt/work/asr_env/.venv/bin/python /opt/work/scripts/asr_eval_soniox.py --wavdir /opt/work/eval/mega_baseline
/opt/work/asr_env/.venv/bin/python /opt/work/scripts/asr_eval_soniox.py --wavdir /opt/work/eval/mega_step1200
echo "FIRST_AB_DONE"

for STEP in 200 400 800; do
  echo "== mega: v2 step $STEP =="
  $PY /opt/work/scripts/synth_eval.py --sentences $MEGA --outdir /opt/work/eval/mega_step$STEP \
      --lora-ckpt "$CKPTDIR/step_$(printf %09d $STEP).ckpt" --lora-config r_32_alpha_16_fast
  /opt/work/asr_env/.venv/bin/python /opt/work/scripts/asr_eval_soniox.py --wavdir /opt/work/eval/mega_step$STEP
done

echo "== MEGA COMPARISON =="
$PY - << "PYEOF"
import json
from pathlib import Path
for d in ["mega_baseline", "mega_step200", "mega_step400", "mega_step800", "mega_step1200"]:
    p = Path(f"/opt/work/eval/{d}/asr_report_soniox.json")
    if p.exists():
        s = json.loads(p.read_text())["summary"]
        tj = Path(f"/opt/work/eval/{d}/timing.jsonl")
        rows = [json.loads(l) for l in tj.read_text().splitlines() if l.strip()]
        ok = [r for r in rows if "rtf" in r]
        dur = round(sum(r["dur_s"] for r in ok), 1) if ok else "-"
        print(f"{d:14s} WER={s['mean_wer']:.3f} CER={s['mean_cer']:.3f} dur={dur}s")
PYEOF
echo "SWEEP_DONE"