#!/usr/bin/env python3
"""P0 基准 · TTS→STT 往返测试。

作用有两个：

  1. 验证 TTS 输出"可懂"——合成得再快，念出来听不懂也没用。
     Piper 的 x_low 是最低音质档，且日志里有 Unknown token 警告，必须确认发出的音是正常的。
  2. 给一个粗略的识别正确率信号（字错率 CER）。

重要局限：这是合成语音，比真人近距离录音干净得多，**会明显高估真实识别率**。
真实的远场识别率必须等麦克风阵列到货后用真人录音测。这里的结果只能当"下限的信号"，
不能当验收数据。
"""

import os
import re
import wave

import numpy as np
import sherpa_onnx

HERE = os.path.dirname(os.path.abspath(__file__))
STT = os.path.join(HERE, "models", "stt")

CASES = [
    "好的我记住了",
    "今天北京晴转多云气温十八到二十六度",
    "明天上午十点有个会议地点在三号会议室",
    "订单编号是七三九一金额一千二百五十元",
    "会议定在二零二六年九月十四日下午三点半",
]


def load_tts(which):
    if which == "piper":
        d = os.path.join(HERE, "models", "tts_piper_xlow")
        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"))
        return sherpa_onnx.OfflineTts(sherpa_onnx.OfflineTtsConfig(
            model=sherpa_onnx.OfflineTtsModelConfig(vits=vits, num_threads=2, provider="cpu"),
            max_num_sentences=1))
    d = os.path.join(HERE, "models", "tts")
    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"))
    return sherpa_onnx.OfflineTts(sherpa_onnx.OfflineTtsConfig(
        model=sherpa_onnx.OfflineTtsModelConfig(vits=vits, num_threads=2, 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_stt():
    return sherpa_onnx.OnlineRecognizer.from_transducer(
        tokens=os.path.join(STT, "tokens.txt"),
        encoder=os.path.join(STT, "encoder-epoch-99-avg-1.int8.onnx"),
        decoder=os.path.join(STT, "decoder-epoch-99-avg-1.int8.onnx"),
        joiner=os.path.join(STT, "joiner-epoch-99-avg-1.int8.onnx"),
        num_threads=2, sample_rate=16000, feature_dim=80,
        decoding_method="greedy_search", enable_endpoint_detection=False)


def transcribe(rec, samples):
    s = rec.create_stream()
    chunk = 1600
    pos = 0
    while pos < len(samples):
        s.accept_waveform(16000, samples[pos:pos + chunk])
        pos += chunk
        while rec.is_ready(s):
            rec.decode_stream(s)
    s.accept_waveform(16000, np.zeros(16000, dtype=np.float32))
    while rec.is_ready(s):
        rec.decode_stream(s)
    s.input_finished()
    while rec.is_ready(s):
        rec.decode_stream(s)
    return rec.get_result(s)


def normalize(t):
    """只保留中文字符做比较，去掉空格与标点。"""
    return re.sub(r"[^一-鿿]", "", t)


def cer(ref, hyp):
    """字错率：编辑距离 / 参考长度。"""
    r, h = normalize(ref), normalize(hyp)
    if not r:
        return float("nan")
    d = np.zeros((len(r) + 1, len(h) + 1), dtype=np.int32)
    d[:, 0] = np.arange(len(r) + 1)
    d[0, :] = np.arange(len(h) + 1)
    for i in range(1, len(r) + 1):
        for j in range(1, len(h) + 1):
            d[i, j] = min(d[i - 1, j] + 1, d[i, j - 1] + 1,
                          d[i - 1, j - 1] + (r[i - 1] != h[j - 1]))
    return d[len(r), len(h)] / len(r)


def main():
    which = os.environ.get("TTS_MODEL", "piper")
    tts = load_tts(which)
    rec = load_stt()
    print(f"TTS={which}  →  STT=zipformer-bilingual-zh-en(int8)\n")

    cers = []
    saved = 0
    outdir = os.path.join(HERE, "roundtrip_wav")
    os.makedirs(outdir, exist_ok=True)
    for i, text in enumerate(CASES):
        audio = tts.generate(text, sid=0, speed=1.0)
        samples = np.array(audio.samples, dtype=np.float32)
        hyp = transcribe(rec, samples)
        c = cer(text, hyp)
        cers.append(c)
        mark = "✅" if c == 0 else "△" if c < 0.2 else "❌"
        print(f"{mark} 原文: {text}")
        print(f"   识别: {hyp}    CER={c:.2f}")
        # 存一份，便于人工听一下 TTS 到底念得对不对
        pcm = (np.clip(samples, -1, 1) * 32767).astype(np.int16)
        with wave.open(os.path.join(outdir, f"{i}_{which}.wav"), "wb") as w:
            w.setnchannels(1)
            w.setsampwidth(2)
            w.setframerate(audio.sample_rate)
            w.writeframes(pcm.tobytes())
        saved += 1

    good = sum(1 for c in cers if c == 0)
    print(f"\n完全正确 {good}/{len(CASES)}，平均 CER={np.mean(cers):.3f}")
    print(f"合成音频已存到 {outdir}/（可人工试听确认音质）")
    print("\n⚠ 合成语音远比真人远场录音干净，此结果不能作为识别率验收数据。")


if __name__ == "__main__":
    main()