"""Mach Boom: text to 24 kHz speech on a CPU.

  python tts.py "Your booking reference is AB-4471-Z." -o out.wav --voice Zoey
"""
import argparse, os, sys, time, wave
import numpy as np
HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
from mach import MachBoom


def pcm16(x):
    return (np.clip(x, -1, 1) * 32767).round().astype("<i2").tobytes()


def write_wav(path, x, sr=24000):
    with wave.open(path, "wb") as w:
        w.setnchannels(1); w.setsampwidth(2); w.setframerate(sr); w.writeframes(pcm16(x))


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("text")
    ap.add_argument("-o", "--out", default="out.wav")
    ap.add_argument("--voice", default="Zoey")
    ap.add_argument("--model", default=os.path.join(HERE, "mach-boom.bin"))
    ap.add_argument("--tokenizer", default=os.path.join(HERE, "tokenizer.json"))
    ap.add_argument("--voices", default=os.path.join(HERE, "voices"))
    ap.add_argument("--threads", type=int, default=1, help="1 to 4: the audio decoder runs alongside the speech model")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--chunk", type=int, default=None, help="decoder chunk in frames (default 16, or 8 above 2 threads)")
    ap.add_argument("--no-frontend", action="store_true", help="skip the English text normalizer")
    ap.add_argument("--stream", action="store_true", help="synthesize piece by piece and print the time to first audio")
    ap.add_argument("--first", type=int, default=None, help="frames in the first piece when streaming (2: first audio "
                    "about 2x sooner, at about 50 dB from the default chunking)")
    ap.add_argument("--pcm", action="store_true", help="stream raw 16-bit 24 kHz mono PCM to stdout, piece by piece")
    a = ap.parse_args()
    tts = MachBoom(a.model, voices_dir=a.voices, tokenizer=a.tokenizer)
    args = dict(seed=a.seed, frontend=not a.no_frontend, threads=a.threads, chunk=a.chunk)
    t0 = time.perf_counter()
    if a.stream or a.pcm:
        pieces, log = [], sys.stderr if a.pcm else sys.stdout
        for x in tts.synth_stream(a.text, a.voice, first=a.first, **args):
            if a.pcm:
                sys.stdout.buffer.write(pcm16(x)); sys.stdout.buffer.flush()
            if not pieces:
                print(f"first audio after {(time.perf_counter() - t0) * 1e3:.0f} ms ({len(x) / 24:.0f} ms of audio)", file=log)
            pieces.append(x)
        wav = np.concatenate(pieces)
    else:
        wav = tts.synth(a.text, a.voice, **args)
    dt = time.perf_counter() - t0
    write_wav(a.out, wav)
    msg = f"{a.out}: {len(wav) / 24000:.2f} s of audio in {dt:.2f} s ({len(wav) / 24000 / dt:.1f}x real time, {tts.k.isa}"
    print(msg + (" + amx)" if tts.amx else ")"), file=sys.stderr if a.pcm else sys.stdout)


if __name__ == "__main__":
    main()