#!/usr/bin/env python3
"""线程数扫描：查清 N5095 上 num_threads 的真实影响。

bench_stt.py 里 2 线程 RTF=0.22、4 线程 RTF=0.87，性能反而掉了 4 倍。
本脚本对 1/2/3/4 线程各重复多轮，排除预热与顺序干扰，确认是否为真实现象。
"""

import os
import statistics
import time
import wave

import numpy as np
import sherpa_onnx

HERE = os.path.dirname(os.path.abspath(__file__))
STT = os.path.join(HERE, "models", "stt")
WAVS = os.path.join(HERE, "wavs")
SR = 16000
CHUNK = int(SR * 0.1)
REPEATS = int(os.environ.get("REPEATS", "5"))


def read_wav(p):
    with wave.open(p, "rb") as w:
        return np.frombuffer(w.readframes(w.getnframes()), dtype=np.int16).astype(np.float32) / 32768.0


def build(nt):
    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=nt, sample_rate=SR, feature_dim=80,
        decoding_method="greedy_search", enable_endpoint_detection=False,
    )


def decode_once(rec, samples):
    s = rec.create_stream()
    t0 = time.perf_counter()
    pos = 0
    while pos < len(samples):
        s.accept_waveform(SR, samples[pos:pos + CHUNK])
        pos += CHUNK
        while rec.is_ready(s):
            rec.decode_stream(s)
    return time.perf_counter() - t0


def main():
    wav = read_wav(os.path.join(WAVS, "stt_0.wav"))
    dur = len(wav) / SR
    print(f"音频 {dur:.2f}s，每档重复 {REPEATS} 轮（含 1 轮预热，不计入）\n")
    print(f"{'线程':>4} {'RTF中位':>9} {'RTF最差':>9} {'各轮RTF':>34}")
    print("-" * 62)
    for nt in (1, 2, 3, 4):
        rec = build(nt)
        decode_once(rec, wav[: SR * 2])  # 预热
        rtfs = []
        for _ in range(REPEATS):
            rtfs.append(decode_once(rec, wav) / dur)
        shown = " ".join(f"{r:.2f}" for r in rtfs)
        print(f"{nt:>4} {statistics.median(rtfs):>9.2f} {max(rtfs):>9.2f} {shown:>34}")
        del rec

    print("\n=== 参照：系统负载 ===")
    print(open("/proc/loadavg").read().strip())


if __name__ == "__main__":
    main()