"""Benchmark MaAI's incremental CPU updates with a full 20s context cache."""
import csv
import json
from pathlib import Path
import time

from smoke_run import HERE, ROOT, np, sf, soxr, torch
from evaluate_can_reply import model


def stats(values):
    v = np.asarray(values, dtype=float)
    return dict(n=len(v), mean_ms=float(v.mean()), median_ms=float(np.median(v)),
                p95_ms=float(np.percentile(v, 95)), p99_ms=float(np.percentile(v, 99)),
                min_ms=float(v.min()), max_ms=float(v.max()))


def main():
    torch.set_num_threads(4)
    torch.set_num_interop_threads(1)
    with (ROOT / "labels/EOT标注/工作数据/eot_labels.csv").open() as f:
        labels = list(csv.DictReader(f))
    limits = {call: max(int(r["decision_ms"]) for r in labels if r["call_id"] == call)
              for call in sorted({r["call_id"] for r in labels})}
    prepared = []
    for call, limit in limits.items():
        pcm, sr = sf.read(ROOT / "data/韩泰机器人抢话/音频" / f"{call}.mp3", dtype="float32", always_2d=True)
        audio = np.ascontiguousarray(soxr.resample(pcm[:limit*sr//1000], sr, 16000, quality="HQ"), dtype=np.float32)
        prepared.append((call, audio))
    measured, load_times = [], []
    for call, audio in prepared:
        began = time.perf_counter()
        maai = model()
        load_times.append((time.perf_counter()-began)*1000)
        maai.CALC_PROCESS_TIME_INTERVAL = 10**9
        current = {}
        original_encode, original_forward = maai.vap.encode_audio, maai.vap.forward

        def encode(*args, **kwargs):
            began = time.perf_counter()
            output = original_encode(*args, **kwargs)
            current["encoder_ms"] = (time.perf_counter()-began)*1000
            return output

        def forward(*args, **kwargs):
            began = time.perf_counter()
            output = original_forward(*args, **kwargs)
            current["vap_ms"] = (time.perf_counter()-began)*1000
            return output

        maai.vap.encode_audio, maai.vap.forward = encode, forward
        count = len(audio)//1600
        for frame in range(count):
            start = frame*1600
            ch1 = np.ascontiguousarray(audio[start:start+1600, 0])
            ch2 = np.ascontiguousarray(audio[start:start+1600, 1])
            current.clear()
            began = time.perf_counter()
            maai.process(ch1, ch2)
            result = maai.result_dict_queue.get_nowait()
            total = (time.perf_counter()-began)*1000
            assert np.isfinite(result["p_now"]).all()
            measured.append(dict(call_id=call, audio_end_sec=(frame+1)*0.1,
                                 full_context=int(frame >= 200), total_ms=total, **current))
        print(f"{call[:8]}: {count} frames, completed", flush=True)
        del maai
    warm = [r for r in measured if r["full_context"]]
    report = dict(method="Incremental process/get-result with 100ms stereo frames, 20s rolling KV cache. Steady-state excludes first 20s of each call. Includes CPC encoding, VAP forward, buffer/queue work; excludes imports, loading, decoding and resampling. No concurrent benchmark jobs.",
                  cpu_threads=4, frame_rate=10, cache_seconds=20,
                  model_load=stats(load_times),
                  all_updates=stats([r["total_ms"] for r in measured]),
                  steady_total=stats([r["total_ms"] for r in warm]),
                  steady_encoder=stats([r["encoder_ms"] for r in warm]),
                  steady_vap=stats([r["vap_ms"] for r in warm]),
                  audio_seconds=len(measured)/10, samples=measured)
    point_rows = json.loads((HERE / "results/work-data-can-reply/summary.json").read_text())
    with (HERE / "results/work-data-can-reply/predictions.csv").open(encoding="utf-8-sig") as f:
        point_replay = list(csv.DictReader(f))
    report["previous_offline_point_replay"] = stats([float(r["replay_inference_ms"]) for r in point_replay])
    report["previous_point_eval_elapsed_seconds"] = point_rows["elapsed_seconds"]
    path = HERE / "results/latency.json"
    path.write_text(json.dumps(report, indent=2)+"\n")
    print(json.dumps({k:v for k,v in report.items() if k != "samples"}, indent=2))


if __name__ == "__main__":
    main()