#!/usr/bin/env python3
"""Parity + accuracy harness for the Audio8-ASR-Infinite ONNX bundle.

  onnx   --prec P --sets jfk jfk_long ...     free-running ORT streaming decode -> results/onnx_P.json
  torch  --draft P --sets jfk jfk_long ...    exact torch fp32 greedy (speculative: ORT tokens of
                                              precision P are only a DRAFT, every token is verified
                                              by a torch forward) -> results/torch.json (+ logits npz)
  shards --sets jfk jfk_long                  fp32 ONNX decoder shards, teacher-forced on the torch
                                              tokens, one shard resident at a time -> results/fp32.json
  report                                      parity / WER / CER table from results/*.json
"""
from __future__ import annotations

import argparse
import json
import os
import re
import time
import unicodedata
from pathlib import Path

import numpy as np

import audio8_infinite_ref as a8i
import audio8_infinite_graph as a8i_onnx

ROOT = Path(__file__).parent
CKPT = Path(os.environ.get("A8I_CKPT", ROOT / "ckpt"))  # HF snapshot of Edge0/Audio8-ASR-Infinite
BUNDLE = Path(os.environ.get("A8I_BUNDLE", ROOT / "bundle"))  # audio8_infinite_export.py --out
RES = Path(os.environ.get("A8I_RESULTS", ROOT / "results"))
RES.mkdir(exist_ok=True)
BENCH = Path(__file__).resolve().parents[1] / "bench" / "audio"
DATA = Path(os.environ.get("A8I_DATA", ROOT / "data"))  # ls_clean/*.f32+.txt, fleurs_cmn_hans_cn/*.f32+.txt


def load_f32(p):
    return np.fromfile(p, dtype="<f4")


def audio_sets(names):
    """name -> list of (utt_id, wav, ref_text|None, lang_token)."""
    out = {}
    for n in names:
        if n == "jfk":
            import soundfile as sf

            x, sr = sf.read(BENCH / "jfk_16k_mono.wav", dtype="float32")
            out[n] = [("jfk", x, "And so, my fellow Americans, ask not what your country can do for you, ask what you can do for your country.", a8i.LANG_EN)]
        elif n == "jfk_long":
            out[n] = [("jfk_long_66s", load_f32(BENCH / "jfk_long_66s.f32"), None, a8i.LANG_EN)]
        elif n == "long_varied":
            out[n] = [("long_varied", load_f32(BENCH / "long_varied.f32"), None, a8i.LANG_EN)]
        elif n in ("ls_clean", "fleurs_zh"):
            d = DATA / ("ls_clean" if n == "ls_clean" else "fleurs_cmn_hans_cn")
            lang = a8i.LANG_EN if n == "ls_clean" else a8i.LANG_ZH
            rows = []
            for f in sorted(d.glob("*.f32")):
                rows.append((f"{n}/{f.stem}", load_f32(f), f.with_suffix(".txt").read_text(encoding="utf-8").strip(), lang))
            out[n] = rows
        else:
            raise SystemExit(f"unknown set {n}")
    return out


# ---------------------------------------------------------------------------------- text metrics
def norm_en(s: str) -> str:
    s = s.lower().replace("’", "'")
    s = re.sub(r"[^a-z0-9' ]+", " ", s)
    return " ".join(s.split())


def norm_zh(s: str) -> str:
    s = unicodedata.normalize("NFKC", s)
    s = "".join(ch for ch in s if not unicodedata.category(ch).startswith(("P", "Z", "S")) and not ch.isspace())
    return s.lower()


def wer(refs, hyps):
    import jiwer

    return jiwer.wer([norm_en(r) for r in refs], [norm_en(h) or "<empty>" for h in hyps])


def cer(refs, hyps):
    import jiwer

    return jiwer.cer([norm_zh(r) for r in refs], [norm_zh(h) or "-" for h in hyps])


# ---------------------------------------------------------------------------------- loaders
def weights_and_dims():
    cfg = json.loads((CKPT / "config.json").read_text())
    return a8i.load_weights(CKPT), a8i.Dims.from_config(cfg)


def tokenizer():
    from tokenizers import Tokenizer

    return Tokenizer.from_file(str(CKPT / "tokenizer.json"))


def ort_backend(prec, threads=None, enc_prec=None):
    sfx = "" if prec == "fp32" else f"_{prec}"
    esfx = sfx if enc_prec is None else ("" if enc_prec == "fp32" else f"_{enc_prec}")
    return a8i_onnx.OrtBackend(BUNDLE / f"audio_encoder{esfx}.onnx", [BUNDLE / f"decoder{sfx}.onnx"], threads=threads)


# ---------------------------------------------------------------------------------- commands
def cmd_onnx(args):
    W, dims = weights_and_dims()
    tok = tokenizer()
    be = ort_backend(args.prec, args.threads, args.enc_prec)
    tag = args.prec if args.enc_prec is None else f"{args.prec}_enc{args.enc_prec}"
    path = RES / f"onnx_{tag}.json"
    res = json.loads(path.read_text(encoding="utf-8")) if path.exists() else {}
    for set_name, rows in audio_sets(args.sets).items():
        for uid, wav, ref, lang in rows:
            if uid in res and not args.keep_logits:
                continue
            s = a8i.Session(be, W, dims, delay_tokens=6, frame_len=4, language=lang, rolling=True,
                            enc_batch=args.enc_batch, keep_logits=args.keep_logits)
            t = time.perf_counter()
            toks = s.run(wav)
            dt = time.perf_counter() - t
            text = a8i.visible_text(tok, toks)
            eot = [round(a8i.eot_probability(v), 4) for v in s.log.vad]
            res[uid] = {"tokens": toks, "text": text, "ref": ref, "seconds": len(wav) / 16000, "wall": dt,
                        "trims": s.trims, "eot": eot}
            if args.keep_logits:
                np.save(RES / f"logits_{tag}_{uid.replace('/', '_')}.npy", np.stack(s.log.logits).astype(np.float32))
            print(f"[{tag}] {uid} {len(wav) / 16000:.1f}s wall={dt:.1f}s rtf={dt / (len(wav) / 16000):.2f} "
                  f"trims={s.trims} :: {text[:120]}", flush=True)
            path.write_text(json.dumps(res, ensure_ascii=False), encoding="utf-8")


def cmd_torch(args):
    import torch

    torch.set_num_threads(args.threads or 16)
    W, dims = weights_and_dims()
    tok = tokenizer()
    ref = a8i.TorchRef(W, dims)
    draft_path = RES / f"onnx_{args.draft}.json"
    drafts = json.loads(draft_path.read_text(encoding="utf-8")) if draft_path.exists() else {}
    out_path = RES / "torch.json"
    res = json.loads(out_path.read_text(encoding="utf-8")) if out_path.exists() else {}
    for set_name, rows in audio_sets(args.sets).items():
        for uid, wav, reftext, lang in rows:
            if uid in res and not args.force:
                continue
            draft = drafts.get(uid, {}).get("tokens", [])
            t = time.perf_counter()
            with torch.no_grad():
                toks, logits, vad, passes = speculative_greedy(ref, W, dims, wav, lang, draft)
            dt = time.perf_counter() - t
            text = a8i.visible_text(tok, toks)
            res[uid] = {"tokens": toks, "text": text, "ref": reftext, "seconds": len(wav) / 16000,
                        "eot": [round(a8i.eot_probability(v), 4) for v in vad], "passes": passes, "wall": dt}
            np.save(RES / f"torchlogits_{uid.replace('/', '_')}.npy", np.stack(logits).astype(np.float32))
            print(f"[torch] {uid} steps={len(toks)} passes={passes} wall={dt:.0f}s :: {text[:120]}", flush=True)
            out_path.write_text(json.dumps(res, ensure_ascii=False), encoding="utf-8")


def speculative_greedy(ref, W, dims, wav, lang, draft):
    """Exact greedy decode with the torch reference; `draft` tokens only decide how many positions a
    single torch forward may verify at once (accepted prefix + the corrected token)."""
    s = a8i.Session(ref, W, dims, delay_tokens=6, frame_len=4, language=lang, rolling=True)
    stream = s.stream_of(wav)
    nwin = s.n_windows(len(stream))
    prompt = [a8i.BOS, lang] + [a8i.STREAM_PAD] * (s.prefill - 2)
    a, b = s.window_bounds(0)
    a0 = s.encode([stream[a:b]], s.prefill)
    audio = []
    k = 1
    while k < nwin:
        kk = list(range(k, min(nwin, k + 24)))
        audio.extend(list(s.encode([stream[s.window_bounds(j)[0]: s.window_bounds(j)[1]] for j in kk], 1)))
        k += len(kk)
    logits_all, vad_all = [], []

    def pick(lg):
        lg = np.asarray(lg, np.float32).reshape(-1).copy()
        lg[a8i.EOS] = -np.inf
        return int(lg.argmax())

    lg, vd = s.decode(a8i.embed_rows(W, prompt) + a0)
    out = [pick(lg)]
    logits_all.append(np.asarray(lg).reshape(-1))
    vad_all.append(np.asarray(vd).reshape(-1))
    passes = 1
    k = 1
    while k < nwin:
        s.maybe_trim(1)
        room = a8i.ROLL_CONTEXT - s.dlen
        m = min(room, nwin - k, 40)
        # inputs for steps k..k+m-1: the token emitted by step j-1 (accepted output, then draft)
        feed = [out[-1]] + [draft[j] if j < len(draft) else a8i.STREAM_PAD for j in range(k, k + m - 1)]
        d0 = s.dlen
        lg, vd = s.decode(a8i.embed_rows(W, feed) + np.stack(audio[k - 1: k - 1 + m]), all_logits=True)
        passes += 1
        lg, vd = np.asarray(lg)[0], np.asarray(vd)[0]
        accepted = m
        for j in range(m):
            t_ = pick(lg[j])
            out.append(t_)
            logits_all.append(lg[j])
            vad_all.append(vd[j].reshape(-1))
            if j + 1 < m and t_ != feed[j + 1]:
                accepted = j + 1
                break
        s.dlen = d0 + accepted
        k += accepted
    return out, logits_all, vad_all, passes


def cmd_shards(args):
    """fp32 ONNX decoder, one shard resident at a time, teacher-forced on torch greedy tokens."""
    import onnxruntime as ort

    W, dims = weights_and_dims()
    tref = json.loads((RES / "torch.json").read_text(encoding="utf-8"))
    enc = a8i_onnx.OrtBackend(BUNDLE / "audio_encoder.onnx", [])
    shards = sorted(BUNDLE.glob("decoder.part*.onnx"))
    out = {}
    for set_name, rows in audio_sets(args.sets).items():
        for uid, wav, _, lang in rows:
            toks = tref[uid]["tokens"]
            s = a8i.Session(enc, W, dims, delay_tokens=6, frame_len=4, language=lang, rolling=True)
            stream = s.stream_of(wav)
            nwin = s.n_windows(len(stream))
            a, b = s.window_bounds(0)
            a0 = s.encode([stream[a:b]], s.prefill)
            audio = []
            k = 1
            while k < nwin:
                kk = list(range(k, min(nwin, k + 24)))
                audio.extend(list(s.encode([stream[s.window_bounds(j)[0]: s.window_bounds(j)[1]] for j in kk], 1)))
                k += len(kk)
            prompt = [a8i.BOS, lang] + [a8i.STREAM_PAD] * (s.prefill - 2)
            inputs = [a8i.embed_rows(W, prompt) + a0] + [a8i.embed_rows(W, [toks[j - 1]]) + audio[j - 1][None] for j in range(1, nwin)]
            logits = None
            for sp in shards:
                be = a8i_onnx.OrtBackend(None, [sp])
                ds = a8i.Session(be, W, dims, delay_tokens=6, frame_len=4, language=lang, rolling=True, ada=s.ada)
                nxt, lgs = [], []
                for x in inputs:
                    ds.maybe_trim(x.shape[0])
                    n = x.shape[0]
                    pos = np.arange(ds.dlen, ds.dlen + n)
                    cos, sin = a8i.rope_cos_sin(pos, dims.t_head_dim, dims.t_theta)
                    bias = np.full((n, ds.L + n), a8i.NEG, np.float32)
                    bias[:, : ds.dlen] = 0.0
                    bias[:, ds.L:] = np.where(np.tril(np.ones((n, n), bool)), 0.0, a8i.NEG)
                    sess = be.decs[0]
                    names_in = {i.name for i in sess.get_inputs()}
                    layers = sorted(int(q.split("_")[-1]) for q in names_in if q.startswith("past_key_"))
                    feed = {("inputs_embeds" if "inputs_embeds" in names_in else "hidden"): x[None].astype(np.float32),
                            "ada_scale": ds.ada, "cos": cos, "sin": sin, "attn_bias": bias}
                    for i in layers:
                        feed[f"past_key_{i}"] = ds.dk[i]
                        feed[f"past_value_{i}"] = ds.dv[i]
                    names = [o.name for o in sess.get_outputs()]
                    r = dict(zip(names, sess.run(names, feed)))
                    for i in layers:
                        ds.dk[i][:, :, ds.dlen: ds.dlen + n] = r[f"key_{i}"]
                        ds.dv[i][:, :, ds.dlen: ds.dlen + n] = r[f"value_{i}"]
                    ds.dlen += n
                    if "hidden" in r:
                        nxt.append(r["hidden"][0])
                    else:
                        lgs.append(r["logits"].reshape(-1))
                del be
                inputs = nxt
                logits = lgs or logits
            tl = np.load(RES / f"torchlogits_{uid.replace('/', '_')}.npy")
            n = min(len(tl), len(logits))
            diffs = [float(np.abs(tl[i] - logits[i]).max()) for i in range(n)]
            am = []
            for i in range(n):
                q = logits[i].copy()
                q[a8i.EOS] = -np.inf
                am.append(int(q.argmax()))
            out[uid] = {"steps": n, "max_abs_logit_diff": max(diffs), "mean_max_abs_logit_diff": float(np.mean(diffs)),
                        "argmax_equal": am == toks[:n], "argmax_mismatches": int(sum(x != y for x, y in zip(am, toks)))}
            print("[fp32-shards]", uid, out[uid], flush=True)
            (RES / "fp32.json").write_text(json.dumps(out, indent=1), encoding="utf-8")


def cmd_report(args):
    torch_res = json.loads((RES / "torch.json").read_text(encoding="utf-8")) if (RES / "torch.json").exists() else {}
    rows = []
    for p in sorted(RES.glob("onnx_*.json")):
        tag = p.stem[5:]
        r = json.loads(p.read_text(encoding="utf-8"))
        line = {"prec": tag}
        for set_name, metric in (("ls_clean", wer), ("fleurs_zh", cer)):
            keys = [k for k in r if k.startswith(set_name + "/")]
            if keys:
                line[set_name] = round(100 * metric([r[k]["ref"] for k in keys], [r[k]["text"] for k in keys]), 2)
                line[set_name + "_n"] = len(keys)
                line[set_name + "_rtf"] = round(sum(r[k]["wall"] for k in keys) / sum(r[k]["seconds"] for k in keys), 3)
        for uid in ("jfk", "jfk_long_66s", "long_varied"):
            if uid in r:
                line[uid + "_rtf"] = round(r[uid]["wall"] / r[uid]["seconds"], 3)
                if uid in torch_res:
                    a, b = r[uid]["tokens"], torch_res[uid]["tokens"]
                    fd = next((i for i in range(min(len(a), len(b))) if a[i] != b[i]), None)
                    line[uid + "_tok_eq"] = (a == b)
                    line[uid + "_first_div"] = fd
                    line[uid + "_agree"] = round(float(np.mean([x == y for x, y in zip(a, b)])), 4)
        # token identity vs torch across all utterances that have a torch run
        common = [k for k in r if k in torch_res]
        if common:
            line["tok_identical_utts"] = f"{sum(r[k]['tokens'] == torch_res[k]['tokens'] for k in common)}/{len(common)}"
        rows.append(line)
    for line in rows:
        print(json.dumps(line, ensure_ascii=False))
    if torch_res:
        for set_name, metric in (("ls_clean", wer), ("fleurs_zh", cer)):
            keys = [k for k in torch_res if k.startswith(set_name + "/")]
            if keys:
                print("torch", set_name, round(100 * metric([torch_res[k]["ref"] for k in keys], [torch_res[k]["text"] for k in keys]), 2), len(keys))


def main():
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)
    a = sub.add_parser("onnx")
    a.add_argument("--prec", required=True)
    a.add_argument("--enc-prec", default=None)
    a.add_argument("--sets", nargs="+", required=True)
    a.add_argument("--threads", type=int, default=None)
    a.add_argument("--enc-batch", type=int, default=16)
    a.add_argument("--keep-logits", action="store_true")
    a.set_defaults(fn=cmd_onnx)
    b = sub.add_parser("torch")
    b.add_argument("--draft", default="int8")
    b.add_argument("--sets", nargs="+", required=True)
    b.add_argument("--threads", type=int, default=None)
    b.add_argument("--force", action="store_true")
    b.set_defaults(fn=cmd_torch)
    c = sub.add_parser("shards")
    c.add_argument("--sets", nargs="+", required=True)
    c.set_defaults(fn=cmd_shards)
    d = sub.add_parser("report")
    d.set_defaults(fn=cmd_report)
    args = ap.parse_args()
    args.fn(args)


if __name__ == "__main__":
    main()