import argparse, json, collections
import soundfile as sf, torch, jiwer
from transformers import AutoModelForCTC, AutoProcessor
SR=16000
def norm(s): return " ".join(str(s).replace("|"," ").split())
ap=argparse.ArgumentParser()
ap.add_argument("--model",required=True)
ap.add_argument("--manifest",default="/root/devhard/devhard_all.jsonl")
a=ap.parse_args()
proc=AutoProcessor.from_pretrained(a.model)
model=AutoModelForCTC.from_pretrained(a.model,torch_dtype=torch.bfloat16).cuda().eval()
rows=[json.loads(l) for l in open(a.manifest,encoding="utf-8")]
by=collections.defaultdict(lambda:{"r":[],"h":[]})
with torch.inference_mode():
    for i,r in enumerate(rows):
        au=sf.read(r["audio"],dtype="float32")[0]
        x=proc(au,sampling_rate=SR,return_tensors="pt")
        x={k:(v.to("cuda",dtype=torch.bfloat16) if v.dtype==torch.float32 else v.to("cuda")) for k,v in x.items()}
        ids=model(**x).logits.argmax(-1).cpu().numpy()
        hyp=norm(proc.batch_decode(ids)[0]); ref=norm(r["text"])
        L=r.get("lang","?"); by[L]["r"].append(ref); by[L]["h"].append(hyp)
        if (i+1)%300==0: print(f"{i+1}/{len(rows)}",flush=True)
scores=[]
for L in sorted(by):
    R=[x for x in by[L]["r"] if x.strip()]; H=[by[L]["h"][j] for j,x in enumerate(by[L]["r"]) if x.strip()]
    wer=jiwer.wer(R,H); cer=jiwer.cer(R,H); comb=0.5*wer+0.5*cer; sc=1-comb
    scores.append(sc)
    print(f"{L}: n={len(R)} WER={wer:.4f} CER={cer:.4f} combine={comb:.4f} score={sc:.4f}")
print(f"MACRO score={sum(scores)/len(scores):.4f}")