#!/bin/bash
# Usage: ckpt_eval.sh <step_XXXXXXXXX.ckpt> [more ckpts...]
# Pauses training, evals baseline (once) + given checkpoints, resumes training.
set -uo pipefail
exec > /opt/work/ckpt_eval2.log 2>&1
CKPTDIR=/opt/work/fish-speech/results/s2pro_egy_lora/checkpoints
NV=/opt/work/asr_env/.venv/lib/python3.12/site-packages/nvidia
PY=/opt/work/fish-speech/.venv/bin/python
echo "=== ckpt_eval $(date) args: $* ==="

echo "== pausing training =="
tmux kill-session -t train 2>/dev/null
pkill -f "fish_speech/train.py" 2>/dev/null
for i in $(seq 1 30); do pgrep -f "fish_speech/train.py" >/dev/null || break; sleep 2; done
sleep 3
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
cd /opt/work/fish-speech

if [ ! -f /opt/work/eval/baseline_true/asr_report_soniox.json ]; then
  echo "== baseline (zero-LoRA) synth =="
  $PY /opt/work/scripts/synth_eval.py --outdir /opt/work/eval/baseline_true
  echo "== baseline ASR =="
  source /opt/work/scripts/soniox.env && \
    /opt/work/asr_env/.venv/bin/python /opt/work/scripts/asr_eval_soniox.py --wavdir /opt/work/eval/baseline_true
fi

for CKPT in "$@"; do
  NAME=$(echo "$CKPT" | grep -oE "step_[0-9]+" | sed "s/step_0*/step/")
  echo "== $NAME synth ($CKPT) =="
  $PY /opt/work/scripts/synth_eval.py --outdir /opt/work/eval/$NAME \
      --lora-ckpt "$CKPTDIR/$CKPT" --lora-config r_32_egy
  echo "== $NAME ASR =="
  source /opt/work/scripts/soniox.env && \
    /opt/work/asr_env/.venv/bin/python /opt/work/scripts/asr_eval_soniox.py --wavdir /opt/work/eval/$NAME
done

echo "== resuming training (auto-resume from latest ckpt) =="
tmux new-session -d -s train "cd /opt/work/fish-speech && export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True TOKENIZERS_PARALLELISM=false && .venv/bin/python fish_speech/train.py --config-name text2semantic_finetune_egy +lora@model.model.lora_config=r_32_egy 2>&1 | tee /opt/work/train.log"
sleep 3; tmux has-session -t train && echo TRAIN_RESUMED || echo TMUX_FAILED

echo "== COMPARISON =="
$PY - << "PYEOF"
import json
from pathlib import Path
for d in sorted(Path("/opt/work/eval").iterdir()):
    rep = d / "asr_report_soniox.json"
    if rep.exists():
        s = json.loads(rep.read_text())["summary"]
        tj = d / "timing.jsonl"
        rtf = "-"
        if tj.exists():
            rows = [json.loads(l) for l in tj.read_text().splitlines() if l.strip()]
            ok = [r for r in rows if "rtf" in r]
            if ok:
                rtf = round(sum(r["rtf"] for r in ok) / len(ok), 2)
        print(f"{d.name:16s} WER={s['mean_wer']:.3f} CER={s['mean_cer']:.3f} RTF={rtf}")
PYEOF
echo "CKPT_EVAL_DONE"