#!/usr/bin/env bash
# I2D (Iterate-to-Differentiate, arXiv 2603.24430): synth N times feeding
# each output back as the next iteration's context. Higher-quality models
# degrade more slowly; the curve over iterations amplifies inter-model
# quality deltas that a single-shot eval would miss.
#
# Outputs:
#   <workdir>/iter_<N>.wav           — synthesized audio per iteration
#   <workdir>/eval_pairs.jsonl       — input to quality_eval (one row per iter)
#   <workdir>/scores.jsonl           — per-iter (speaker_cosine, wer, peak/rms)
#   <workdir>/curve.tsv              — pretty-printed degradation curve
#
# Reference for speaker similarity stays the original ref_wav; reference
# for WER stays the original SYNTH_TEXT. Both are constant across iters,
# so any drift you see is the model degrading, not the target moving.
#
# Usage:
#   scripts/i2d_loop.sh \
#       --text "What you want each iteration to say." \
#       --ref-wav /tmp/voice/picked.wav \
#       --ref-text "Transcript of the reference clip." \
#       --iterations 5 \
#       --workdir /tmp/i2d_run \
#       --wavlm-sv /tmp/wavlm_sv.safetensors

set -euo pipefail

TEXT=""
REF_WAV=""
REF_TEXT=""
ITERATIONS=5
WORKDIR=""
WAVLM_SV="/tmp/wavlm_sv.safetensors"
SPEAKER=0
SEED=42

while [[ $# -gt 0 ]]; do
    case "$1" in
        --text) TEXT="$2"; shift 2 ;;
        --ref-wav) REF_WAV="$2"; shift 2 ;;
        --ref-text) REF_TEXT="$2"; shift 2 ;;
        --iterations) ITERATIONS="$2"; shift 2 ;;
        --workdir) WORKDIR="$2"; shift 2 ;;
        --wavlm-sv) WAVLM_SV="$2"; shift 2 ;;
        --speaker) SPEAKER="$2"; shift 2 ;;
        --seed) SEED="$2"; shift 2 ;;
        -h|--help)
            sed -n '2,28p' "$0"
            exit 0
            ;;
        *) echo "unknown arg: $1" >&2; exit 1 ;;
    esac
done

if [[ -z "$TEXT" || -z "$REF_WAV" || -z "$REF_TEXT" || -z "$WORKDIR" ]]; then
    echo "usage: $0 --text <t> --ref-wav <w> --ref-text <r> --workdir <d> [--iterations N] [--wavlm-sv <f>]" >&2
    exit 1
fi
if [[ ! -f "$REF_WAV" ]]; then
    echo "ref-wav not found: $REF_WAV" >&2; exit 1
fi
if [[ ! -f "$WAVLM_SV" ]]; then
    echo "wavlm-sv not found: $WAVLM_SV" >&2; exit 1
fi

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
WORKSPACE_DIR="$(cd "$SCRIPT_DIR/../../../.." && pwd)"
GEN_BIN="$WORKSPACE_DIR/target/release/examples/generate"
EVAL_BIN="$WORKSPACE_DIR/target/release/examples/quality_eval"
for bin in "$GEN_BIN" "$EVAL_BIN"; do
    if [[ ! -x "$bin" ]]; then
        echo "binary missing: $bin" >&2
        echo "build with: cargo build -p rtx-csm --release --example generate --example quality_eval" >&2
        exit 1
    fi
done

mkdir -p "$WORKDIR"
PAIRS="$WORKDIR/eval_pairs.jsonl"
: > "$PAIRS"

# Iter 0 uses the original ref_wav as context.
# Iter N>0 uses iter_(N-1).wav as context, with TEXT as the context text
# (since the previous iteration synthesized TEXT).
PREV_WAV="$REF_WAV"
PREV_TEXT="$REF_TEXT"

for ((i=0; i<ITERATIONS; i++)); do
    OUT_WAV="$WORKDIR/iter_${i}.wav"
    echo "→ iter $i: synth → $(basename "$OUT_WAV")  (ctx=$(basename "$PREV_WAV"))"
    "$GEN_BIN" \
        --text "$TEXT" \
        --speaker "$SPEAKER" \
        --context-wav "$PREV_WAV" \
        --context-text "$PREV_TEXT" \
        --context-speaker 1 \
        --seed "$SEED" \
        --out "$OUT_WAV" 2>&1 | grep -E "generated" | tail -1 || true

    if [[ ! -f "$OUT_WAV" ]]; then
        echo "iter $i produced no output — aborting" >&2
        exit 2
    fi

    # Append eval row: ref always = the *original* reference so speaker
    # cosine measures drift from the source voice; WER target = TEXT.
    jq -nc --arg ref "$REF_WAV" --arg gen "$OUT_WAV" \
        --arg t "$TEXT" --argjson iter "$i" \
        '{ref_wav: $ref, gen_wav: $gen, ref_text: $t, iter: $iter}' >> "$PAIRS"

    PREV_WAV="$OUT_WAV"
    PREV_TEXT="$TEXT"
done

echo
echo "→ scoring all $ITERATIONS iterations"
"$EVAL_BIN" \
    --in "$PAIRS" \
    --out "$WORKDIR/scores.jsonl" \
    --wavlm-sv "$WAVLM_SV" 2>&1 | tail -5

echo
echo "=== degradation curve ==="
{
    echo "iter	speaker_cosine	wer	peak_db	rms_db"
    jq -r '[.iter, .metrics.speaker_cosine, .metrics.wer, .metrics.gen_peak_db, .metrics.gen_rms_db] | @tsv' \
        "$WORKDIR/scores.jsonl"
} | tee "$WORKDIR/curve.tsv" | column -t -s $'\t'

echo
echo "✓ wrote $WORKDIR/curve.tsv"
echo "  audio in $WORKDIR/iter_*.wav"