"""Evaluate an omniASR CTC model (+ optional fine-tuned checkpoint) with the duplex decoding (greedy CTC + Latin-script mask).
~/venv_omni/bin/python eval_asr.py --tag NAME [--card omniASR_CTC_1B_v2] [--ckpt /scratch/asr/runs/A/best.pt] [--sets duplex,realcs,synthtest,ace]
  --hyp_only MANIFEST : just transcribe a manifest (used for the synthetic-data CER filter) -> MANIFEST.hyp.jsonl
Writes /scratch/asr/eval/<tag>.json (scores) and <tag>_hyps.jsonl."""
import argparse, json, os, re, unicodedata, time, numpy as np, torch, torch.nn.functional as F
import soundfile as sf
from scipy.signal import resample_poly
from fairseq2.models.hub import load_model
from fairseq2.data.tokenizers.hub import load_tokenizer
from fairseq2.nn.batch_layout import BatchLayout
import jiwer
ap = argparse.ArgumentParser()
ap.add_argument("--tag", default="x"); ap.add_argument("--card", default="omniASR_CTC_1B_v2"); ap.add_argument("--ckpt", default="")
ap.add_argument("--sets", default="duplex,realcs,synthtest,ace"); ap.add_argument("--hyp_only", default=""); ap.add_argument("--latin", type=int, default=1)
a = ap.parse_args()
dev = torch.device("cuda")
model = load_model(a.card, device=dev, dtype=torch.bfloat16).eval()
if a.ckpt:
    ck = torch.load(a.ckpt, map_location="cuda", weights_only=False); print("load", a.ckpt, model.load_state_dict(ck["model"]), ck.get("step"), ck.get("dev"))
tok = load_tokenizer(a.card); dec = tok.create_decoder(skip_special_tokens=True)
V = model.final_proj.weight.shape[0]
def _ok(sx):
    for c in sx:
        cat = unicodedata.category(c)
        if cat[0] == "M":
            if not ("\u0300" <= c <= "\u036f"): return False
        elif c.isalpha() and "LATIN" not in unicodedata.name(c, ""): return False
    return True
BAD = torch.tensor([i for i in range(V) if not _ok(str(dec(torch.tensor([i]))))], dtype=torch.long, device=dev) if a.latin else None
print("latin mask blocks", 0 if BAD is None else len(BAD), "of", V)

def load(r):
    if r.get("wav"):
        y, sr = sf.read(r["wav"], dtype="float32")
        if y.ndim > 1: y = y.mean(1)
        return y if sr == 16000 else resample_poly(y, 16000, sr).astype(np.float32)
    m = np.memmap(r["pcm"], dtype=np.int16, mode="r"); y = m[r["off24"]: r["off24"] + r["n24"]].astype(np.float32) / 32768
    return resample_poly(y, 2, 3).astype(np.float32)

@torch.no_grad()
def transcribe(items, bs=8):
    out, lat = [], []
    order = sorted(range(len(items)), key=lambda i: items[i].get("dur", 0)); res = [None] * len(items)
    for s in range(0, len(order), bs):
        idx = order[s: s + bs]; ys = [torch.from_numpy(load(items[i])) for i in idx]
        ys = [F.layer_norm(y, y.shape) for y in ys]; L = [len(y) for y in ys]
        X = torch.zeros(len(ys), max(L)); [X[k, : len(y)].copy_(y) for k, y in enumerate(ys)]
        X = X.to(dev, torch.bfloat16); bl = BatchLayout(X.shape, seq_lens=L, device=dev)
        lg, blo = model(X, bl)
        if BAD is not None: lg[..., BAD] = -float("inf")
        pr = lg.argmax(-1)
        for k, i in enumerate(idx):
            sq = pr[k, : blo.seq_lens[k]]; m = torch.ones_like(sq, dtype=torch.bool); m[1:] = sq[1:] != sq[:-1]
            res[i] = str(dec(sq[m]))
    return res

def n(s):
    s = unicodedata.normalize("NFC", s.lower()); s = re.sub(r"[^\w\s']", " ", s); return re.sub(r"\s+", " ", s).strip()
def fold(s): return "".join(c for c in unicodedata.normalize("NFD", s) if unicodedata.category(c) != "Mn")
def score(refs, hyps):
    R = [n(r) or "x" for r in refs]; H = [n(h) or "x" for h in hyps]
    return dict(n=len(R), wer=round(jiwer.wer(R, H), 4), cer=round(jiwer.cer(R, H), 4),
                cer_nodiac=round(jiwer.cer([fold(r) for r in R], [fold(h) for h in H]), 4))

if a.hyp_only:
    items = [json.loads(l) for l in open(a.hyp_only)]; hy = transcribe(items, 16)
    with open(a.hyp_only + ".hyp.jsonl", "w") as f:
        for r, h in zip(items, hy):
            f.write(json.dumps(dict(id=r["id"], hyp=h, cer=round(jiwer.cer(n(r["text"]) or "x", n(h) or "x"), 4)), ensure_ascii=False) + "\n")
    print("HYP_DONE", len(items)); raise SystemExit

os.makedirs("/scratch/asr/eval", exist_ok=True); res = {}; allh = []
sets = a.sets.split(",")
if "duplex" in sets:   # the duplex REPORT set: WAXAL test x40 ibo/yor/hau/pcm, FLEURS val eng x29, 40 concatenated cs clips
    R = [json.loads(l) for l in open("/scratch/duplex/asr_eval/list.jsonl")]; H = transcribe(R)
    for lg in ["eng", "pcm", "ibo", "yor", "hau", "cs"]:
        ix = [i for i, r in enumerate(R) if r["lang"] == lg]; res[f"duplex_{lg}"] = score([R[i]["text"] for i in ix], [H[i] for i in ix])
    allh += [dict(set="duplex", id=r["id"], ref=r["text"], hyp=h) for r, h in zip(R, H)]
if "realcs" in sets:   # real speech with English inside Igbo/Yoruba/Hausa sentences (held out)
    R = [json.loads(l) for l in open("/scratch/asr/data/realcs_test.jsonl")]; H = transcribe(R)
    res["realcs_all"] = score([r["text"] for r in R], H)
    for k in ["natural_conversational", "insertional"]:
        ix = [i for i, r in enumerate(R) if r["cs_kind"] == k]; res[f"realcs_{k}"] = score([R[i]["text"] for i in ix], [H[i] for i in ix])
    hit = tot = 0   # English-word recall: share of the English words in the reference that appear in the hypothesis
    for r, h in zip(R, H):
        hw = n(h).split()
        for w in r["cs_eng"]:
            tot += 1; hit += w in hw
    res["realcs_eng_word_recall"] = round(hit / max(1, tot), 4)
    allh += [dict(set="realcs", id=r["id"], ref=r["text"], hyp=h) for r, h in zip(R, H)]
if "synthtest" in sets and os.path.exists("/scratch/asr/synth/test.jsonl"):   # held-out synthetic cs texts (TTS domain)
    R = [json.loads(l) for l in open("/scratch/asr/synth/test.jsonl")]; H = transcribe(R); res["synthtest"] = score([r["text"] for r in R], H)
    allh += [dict(set="synthtest", id=r["id"], ref=r["text"], hyp=h) for r, h in zip(R, H)]
if "ace" in sets and os.path.exists("/scratch/asr/ace/list.jsonl"):   # ACE's own recordings, when available
    R = [json.loads(l) for l in open("/scratch/asr/ace/list.jsonl")]; H = transcribe(R, 1); res["ace"] = score([r["text"] for r in R], H)
    allh += [dict(set="ace", id=r["id"], ref=r["text"], hyp=h) for r, h in zip(R, H)]
json.dump(dict(tag=a.tag, card=a.card, ckpt=a.ckpt, scores=res), open(f"/scratch/asr/eval/{a.tag}.json", "w"), indent=1)
with open(f"/scratch/asr/eval/{a.tag}_hyps.jsonl", "w") as f:
    for h in allh: f.write(json.dumps(h, ensure_ascii=False) + "\n")
print(json.dumps(res, indent=1)); print("EVAL_DONE", a.tag)