"""Run the PiVo Python reference and dump intermediates for JS parity tests."""
import json, sys
from pathlib import Path
import numpy as np

ROOT = Path(__file__).resolve().parents[3]  # the workspace that holds model/ and ref/
sys.path.insert(0, str(ROOT / "model"))
import pivo.text as T
import pivo.runtime as R
from pivo.audio import write_wav

# The app uses the frontend's opt-in letter-spelling fallback for unknown words.
R.phonemize = lambda text, overrides=None: T.phonemize(text, strict=False, overrides=overrides)

out = ROOT / "ref"
out.mkdir(exist_ok=True)
sentences = json.loads((Path(__file__).resolve().parents[1] / "sentences.json").read_text())
rt = R.Runtime(ROOT / "model/onnx", threads=1)

# Capture the decoder features fed to the vocoder.
orig_run = {}
captured = {}
for name, sess in rt.sessions.items():
    class Wrap:
        def __init__(self, s, n): self.s, self.n = s, n
        def run(self, outs, feeds):
            r = self.s.run(outs, feeds)
            captured.setdefault(self.n, []).append((feeds, r))
            return r
    rt.sessions[name] = Wrap(sess, name)

frontend = []
for i, s in enumerate(sentences):
    fr = T.phonemize(s, strict=False)
    frontend.append({"text": s, "normalized": fr.normalized_text, "symbols": list(fr.symbols),
                     "ids": list(fr.token_ids), "unknown": list(fr.unknown_words)})
    for voice in ["Pi", "Vo"]:
        captured.clear()
        audio = rt.synthesize(s, voice=voice)
        enc_feeds, enc_out = captured["encoder"][0]
        feats = np.concatenate([r[0][0] for _, r in captured["decoder"]], 0)
        stem = f"s{i}_{voice}"
        np.save(out / f"{stem}_encoded.npy", enc_out[0][0])
        np.save(out / f"{stem}_logdur.npy", enc_out[1][0])
        np.save(out / f"{stem}_features.npy", feats)
        np.save(out / f"{stem}_audio.npy", audio.astype(np.float32))
        write_wav(out / f"{stem}.wav", audio)
        print(stem, len(fr.token_ids), "tokens", feats.shape[0], "frames", f"{len(audio)/16000:.2f}s")

# Extra frontend-only cases.
extra = ["", "a", "Mr. and Mrs. Jones vs. the St. Louis Prof.", "I'll be there at 9:30, ok?", "x-ray, re-enter, e-mail",
         "THE END", "USA FBI ABCDEF", "Pi's voice isn't Vo's.", "5th 1st 2nd 3rd 100th 11th 12th 13th 20th 21st 101st",
         "2.5 0.75 -3.25 1,000 99,999,999 1000000000000000",  "$0.01 $1 $1.00 £1.01 €100.5", "rock 'n' roll", "Café naïve résumé",
         "ellipsis… and “smart quotes” and ‘single’ — dash – en", "hello;world:foo!bar?baz.", "x"]
for s in extra:
    try:
        fr = T.phonemize(s, strict=False)
        frontend.append({"text": s, "normalized": fr.normalized_text, "symbols": list(fr.symbols),
                         "ids": list(fr.token_ids), "unknown": list(fr.unknown_words)})
    except Exception as e:
        frontend.append({"text": s, "error": f"{type(e).__name__}: {e}"})
(out / "frontend.json").write_text(json.dumps(frontend, indent=1, ensure_ascii=False))