"""Write pivo_runner's inputs for one or more texts: the author's frontend turns each text into phone ids.

For text k (0, 1, ...) in the given order, <out>/<k as 3 digits>_tokens.i32 holds the phone ids (int32, little endian)
and <k>_voice.i32 one int32 (0 = Pi, 1 = Vo); index.tsv lists them. The runner reads them with --fixtures <out>
--count <number of texts>.

Run from the repository root (the host's packages: pip install -r host/requirements.txt):
    python runner/make_inputs.py --out inputs "NASA is ready for the next mission."
    python runner/make_inputs.py --out inputs --voice Vo "Hello, my name is Vo." "Good morning."
"""
import argparse
import os
import sys

import numpy as np

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(ROOT, "host"))

from pivo.text import SPECIAL, VOICE_IDS, phonemize  # noqa: E402  (the author's frontend, host/pivo/)

MAX_TOKENS = 512  # the largest encoder graph


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("texts", nargs="+")
    ap.add_argument("--out", required=True)
    ap.add_argument("--voice", choices=sorted(VOICE_IDS), default="Pi")
    a = ap.parse_args()
    os.makedirs(a.out, exist_ok=True)
    rows = ["# id3\tvoice\ttokens\ttext"]
    for k, text in enumerate(a.texts):
        front = phonemize(text)
        if not any(s not in SPECIAL for s in front.symbols):
            raise SystemExit(f"text {k}: no spoken text")
        if len(front.token_ids) > MAX_TOKENS:
            raise SystemExit(f"text {k}: {len(front.token_ids)} tokens > {MAX_TOKENS}; split it")
        id3 = f"{k:03d}"
        np.array(front.token_ids, "<i4").tofile(os.path.join(a.out, f"{id3}_tokens.i32"))
        np.array([VOICE_IDS[a.voice]], "<i4").tofile(os.path.join(a.out, f"{id3}_voice.i32"))
        rows.append(f"{id3}\t{a.voice}\t{len(front.token_ids)}\t{text}")
    with open(os.path.join(a.out, "index.tsv"), "w", encoding="utf-8") as f:
        f.write("\n".join(rows) + "\n")
    print(f"{len(a.texts)} inputs in {a.out}: run pivo_runner with --fixtures {a.out} --count {len(a.texts)}")


if __name__ == "__main__":
    main()