"""Validate haas_coreml.py end to end.
1) wav33 (all <= 15 s, one window): token match with results/wav33/pynemo_fp32.json, RTF.
2) long audio (wav33 concatenated with 0.4 s silence, ~45-55 s): CER against the joined per-clip NeMo texts for
   CoreML chunked, fp32 ONNX with the same chunking, and fp32 ONNX on the whole utterance (no chunking).
Writes results/coreml_e2e.json."""
import json, os, sys, time, unicodedata
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import numpy as np
from rapidfuzz.distance import Levenshtein
from haas_coreml import ONNX_DIR, SR, SUB, WINDOW, Haas, chunks, load_wav, sessions

# results/wav33/pynemo_fp32.json and wav33.txt (the eval audio list) are not distributed; run from a dir that has them.
h = Haas()
enc_onnx, _ = sessions(os.path.join(ONNX_DIR, "encoder.onnx"), os.path.join(ONNX_DIR, "decoder_joint.onnx"), 4)


def onnx_encode(mel, L, chunked):
    def run(x, n):
        y, _ = enc_onnx.run(None, {"audio_signal": x, "length": np.array([n], np.int64)})
        return y
    if not chunked:
        return run(np.ascontiguousarray(mel[:, :, :L]), L)
    outs = []
    for start, k0, k1 in chunks(L):
        n = min(WINDOW, L - start)
        x = np.zeros((1, 128, WINDOW), np.float32); x[:, :, :n] = mel[:, :, start:start + n]
        outs.append(run(x, n)[:, :, (k0 - start) // SUB:(k1 - start + SUB - 1) // SUB])
    return np.concatenate(outs, axis=2)


def norm(t):
    t = unicodedata.normalize("NFKC", t).lower()
    return "".join(c for c in t if unicodedata.category(c)[0] not in "PZS")


def decode(enc):
    return [int(i) for i in __import__("haas_coreml").greedy_tdt(h.dec, enc, enc.shape[2])]


def main():
    ref = json.load(open("results/wav33/pynemo_fp32.json"))
    paths = open("wav33.txt").read().split()
    items = [(os.path.splitext(os.path.basename(p))[0], load_wav(p)) for p in paths]
    h.transcribe(items[0][1])  # warmup / ANE load
    out = {"wav33": {}, "long": {}}
    match, t_all, dur = 0, 0.0, 0.0
    for k, a in items:
        t0 = time.time(); ids, text = h.transcribe(a); t_all += time.time() - t0; dur += len(a) / SR
        match += ids == ref[k]["ids"]
        out["wav33"][k] = {"ids": ids, "text": text, "match": ids == ref[k]["ids"]}
        if ids != ref[k]["ids"]:
            print("DIFF", k, "\n  ref:", ref[k]["text"], "\n  hyp:", text)
    print(json.dumps({"wav33_match": f"{match}/{len(items)}", "e2e_rtf": round(t_all / dur, 4), "audio_s": round(dur, 1)}))

    # long audio: groups of consecutive clips
    sil = np.zeros(int(0.4 * SR), np.float32)
    groups, cur = [], []
    for k, a in items:
        cur.append((k, a))
        if sum(len(x) for _, x in cur) / SR > 45:
            groups.append(cur); cur = []
    if cur:
        groups.append(cur)
    tot = {k: [0, 0] for k in ["coreml_chunked", "onnx_chunked", "onnx_full"]}
    for gi, g in enumerate(groups):
        a = np.concatenate([np.concatenate([x, sil]) for _, x in g])
        ref_text = "".join(ref[k]["text"] for k, _ in g)
        mel, L = h.front(a)
        t0 = time.time(); _, c_text = h.transcribe(a); dt = time.time() - t0
        texts = {"coreml_chunked": c_text, "onnx_chunked": h.tok.decode(decode(onnx_encode(mel, L, True))),
                 "onnx_full": h.tok.decode(decode(onnx_encode(mel, L, False)))}
        row = {"keys": [k for k, _ in g], "audio_s": round(len(a) / SR, 1), "n_chunks": len(chunks(L)),
               "rtf": round(dt / (len(a) / SR), 4), "ref_text": ref_text, "texts": texts}
        for name, t in texts.items():
            e = Levenshtein.distance(norm(t), norm(ref_text)); row[f"{name}_cer"] = round(100 * e / len(norm(ref_text)), 2)
            tot[name][0] += e; tot[name][1] += len(norm(ref_text))
        out["long"][f"group{gi}"] = row
        print(json.dumps({k: v for k, v in row.items() if "text" not in k}))
    summ = {k: f"{100 * v[0] / v[1]:.2f}% ({v[0]}/{v[1]})" for k, v in tot.items()}
    print(json.dumps({"long_cer_vs_joined_clip_refs": summ}))
    out["summary"] = summ
    json.dump(out, open("results/coreml_e2e.json", "w"), ensure_ascii=False, indent=1)


if __name__ == "__main__":
    main()