"""End-to-end ONNX inference: text -> acoustic.onnx -> mel -> bigvgan.onnx -> wav."""

import argparse
import os
import sys

import numpy as np
import onnxruntime as ort
import soundfile as sf

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from cn_frontend import text_to_sequence


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--text", required=True)
    parser.add_argument("--acoustic", default="export/acoustic_female.onnx")
    parser.add_argument("--vocoder", default="export/bigvgan_base.onnx")
    parser.add_argument("--output", default="export/onnx_demo.wav")
    parser.add_argument("--max_text", type=int, default=256)
    parser.add_argument("--max_mel", type=int, default=2048)
    parser.add_argument("--noise_scale", type=float, default=0.3)
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()

    from infer_board import split_sentences, text_to_inputs, tail_stretch
    ac = ort.InferenceSession(args.acoustic, providers=["CPUExecutionProvider"])
    vc = ort.InferenceSession(args.vocoder, providers=["CPUExecutionProvider"])

    pieces = []
    for s in split_sentences(args.text):
        x, x_lengths, noise_z = text_to_inputs(s, args.noise_scale, args.seed)
        mel, y_lengths = ac.run(None, {
            "x": x, "x_lengths": x_lengths, "noise_z": noise_z,
        })
        T = int(y_lengths[0])
        mel = mel[:, :, :T]
        mel = tail_stretch(mel)
        print(f"  sentence mel: {mel.shape} (y_lengths={T})")
        wav = vc.run(None, {"mel": mel})[0]  # [1,1,T*256]
        pieces.append(wav[0, 0])
    wav = np.concatenate(pieces) if len(pieces) > 1 else pieces[0]
    wav = wav / (np.abs(wav).max() + 1e-8) * 0.95
    sf.write(args.output, wav, 24000)
    print(f"saved: {args.output} ({len(wav)/24000:.2f}s)")


if __name__ == "__main__":
    main()