#!/usr/bin/env python3
"""P0 基准 · 本地 TTS 模型横向对比（N5095）。

背景：vits-zh-ll 实测 RTF≈1.7，慢于实时，且 VITS 非自回归——
回调只在合成结束后触发一次，首音延迟 == 总合成耗时。
本脚本用同一组句子对比多个候选模型，找是否有能跑到 RTF<0.5 的。

注意 max_num_sentences=1 会把多句输入截断为第一句（见 bench_tts.py 的坑），
所以这里每句单独调用——这也正是真实系统的做法（按语义短句合成）。
"""

import os
import statistics
import time

import numpy as np
import sherpa_onnx

HERE = os.path.dirname(os.path.abspath(__file__))

SHORT = "好的，我记住了。"
MED = "今天北京晴转多云，气温十八到二十六度。"
LONG = "我查了一下，明天上午十点有个会议，地点在三号会议室。"
NUM = "订单编号是七三九一，金额一千二百五十元六角。"

TESTS = [("短句", SHORT), ("中句", MED), ("长句", LONG), ("数字", NUM)]


def load_vits_zh_ll(nt):
    d = os.path.join(HERE, "models", "tts")
    return sherpa_onnx.OfflineTts(sherpa_onnx.OfflineTtsConfig(
        model=sherpa_onnx.OfflineTtsModelConfig(
            vits=sherpa_onnx.OfflineTtsVitsModelConfig(
                model=os.path.join(d, "model.onnx"),
                lexicon=os.path.join(d, "lexicon.txt"),
                tokens=os.path.join(d, "tokens.txt"),
                dict_dir=os.path.join(d, "dict"),
            ),
            num_threads=nt, provider="cpu"),
        rule_fsts=",".join(os.path.join(d, f) for f in
                           ("date.fst", "number.fst", "phone.fst", "new_heteronym.fst")),
        max_num_sentences=1))


def load_piper_xlow(nt):
    d = os.path.join(HERE, "models", "tts_piper_xlow")
    return sherpa_onnx.OfflineTts(sherpa_onnx.OfflineTtsConfig(
        model=sherpa_onnx.OfflineTtsModelConfig(
            vits=sherpa_onnx.OfflineTtsVitsModelConfig(
                model=os.path.join(d, "zh_CN-huayan-x_low.onnx"),
                lexicon=os.path.join(d, "lexicon.txt"),
                tokens=os.path.join(d, "tokens.txt"),
            ),
            num_threads=nt, provider="cpu"),
        max_num_sentences=1))


MODELS = [
    ("vits-zh-ll", load_vits_zh_ll, "VITS 中文，115MB"),
    ("piper-huayan-x_low", load_piper_xlow, "Piper 中文，20MB"),
]


def measure(tts, text):
    calls = [0]

    def cb(*args):
        calls[0] += 1
        return 0

    t0 = time.perf_counter()
    audio = tts.generate(text, sid=0, speed=1.0, callback=cb)
    total_ms = (time.perf_counter() - t0) * 1000
    dur = len(audio.samples) / audio.sample_rate
    return total_ms, dur, calls[0]


def main():
    nt = int(os.environ.get("THREADS", "2"))
    print(f"N5095 / num_threads={nt} / 每句单独合成\n")

    report = {}
    for name, loader, desc in MODELS:
        try:
            tts = loader(nt)
        except Exception as e:
            print(f"=== {name} ({desc}) ===\n  加载失败: {type(e).__name__}: {e}\n")
            continue
        tts.generate("预热", sid=0, speed=1.0)
        print(f"=== {name}  ({desc}) ===")
        rows = []
        for label, text in TESTS:
            total_ms, dur, ncalls = measure(tts, text)
            rtf = (total_ms / 1000) / dur
            rows.append((label, total_ms, dur, rtf, ncalls))
            print(f"  {label:<5} {len(text):>2}字 → 音频{dur:5.2f}s  合成{total_ms:7.0f}ms  "
                  f"RTF={rtf:5.2f}  回调{ncalls}次")
        rtfs = [r[3] for r in rows]
        avg = statistics.mean(rtfs)
        print(f"  → RTF 均值={avg:.2f}  最差={max(rtfs):.2f}  "
              f"判定: {'✅ 可实时' if avg < 0.5 else '⚠ 慢于实时' if avg >= 1 else '△ 勉强'}\n")
        report[name] = {"desc": desc, "rows": rows, "rtf_avg": avg, "rtf_max": max(rtfs)}
        del tts

    # 播放侧模拟：连续说 5 句时，下一句能否在上一句播完前备好
    print("=== 连续播放模拟（5 句，串行合成）===")
    loaders = {m[0]: m[1] for m in MODELS}
    for name in loaders:
        if name not in report:
            continue
        tts = loaders[name](nt)
        sents = ["我查了一下，明天上午十点有个会议。", "地点在三号会议室。",
                 "另外提醒你，后天是妈妈的生日。", "要不要我提前帮你订一束花？",
                 "还有你上周说要买的那本书现在有货了。"]
        t_gen = 0.0
        play_end = 0.0
        gaps = []
        for i, s in enumerate(sents):
            total_ms, dur, _ = measure(tts, s)
            first_ready = t_gen + total_ms / 1000
            play_start = max(play_end, first_ready)
            gaps.append(max(0.0, play_start - play_end))
            play_end = play_start + dur
            t_gen += total_ms / 1000
        print(f"  {name:<22} 各句空档: {[f'{g:.1f}s' for g in gaps]}  "
              f"最大={max(gaps):.2f}s  总播完={play_end:.1f}s")
        report[name]["gaps"] = gaps
        del tts
    print()

    import json
    out = os.path.join(HERE, "results_tts_compare.json")
    with open(out, "w") as f:
        json.dump(report, f, ensure_ascii=False, indent=2)
    print(f"结果已写入 {out}")


if __name__ == "__main__":
    main()