"""Offline X-ASR (zh-en zipformer transducer, 2026-06-03) scorer — accurate code-mix
assessment. Reuses the model's own raw-ort decode (int64-patched) + OpenCC T2S norm.
Run in moss-nano-venv (onnxruntime 1.27 + kaldi_native_fbank + opencc + soundfile/librosa)."""
import os, re, sys
MDIR="/home/luigi/jetson-tts/models/xasr-offline/sherpa-onnx-x-asr-zipformer-transducer-zh-en-2026-06-03"
_cwd=os.getcwd(); os.chdir(MDIR); sys.path.insert(0, MDIR)
import importlib.util
_spec=importlib.util.spec_from_file_location("xasr_test", f"{MDIR}/test_onnx.py")
T=importlib.util.module_from_spec(_spec); _spec.loader.exec_module(T)
os.chdir(_cwd)
import opencc, cn2an
_ZD={"0":"零","1":"一","2":"二","3":"三","4":"四","5":"五","6":"六","7":"七","8":"八","9":"九"}
def _numnorm(t):
    def r(m):
        d=m.group(0)
        if len(d)>=5: return "".join(_ZD[c] for c in d)
        try: return cn2an.an2cn(d,"low")
        except: return "".join(_ZD[c] for c in d)
    return re.sub(r"\d+",r,t)
_t2s=opencc.OpenCC('t2s'); _model=None; _id2tok=None
def _ensure():
    global _model,_id2tok
    if _model is None:
        cwd=os.getcwd(); os.chdir(MDIR)
        _id2tok=T.load_tokens("./tokens.txt"); _model=T.load_model(use_int8=False)
        os.chdir(cwd)
    return _model,_id2tok
def asr(wav):
    import soundfile as sf, librosa, numpy as np
    m,id2tok=_ensure()
    samples,sr=sf.read(wav,dtype="float32")
    if samples.ndim>1: samples=samples.mean(1)
    if sr!=16000: samples=librosa.resample(samples,orig_sr=sr,target_sr=16000,res_type="soxr_hq")
    if len(samples)<64000: samples=np.concatenate([samples,np.zeros(64000-len(samples),np.float32)])  # >=4s pad: short clips break encoder convs
    feats=T.compute_feat(samples=samples.astype(np.float32), sample_rate=16000)
    blank=0; hyp=[blank]*m.context_size
    dout=m.run_decoder(hyp); enc=m.run_encoder(feats[None])
    for k in range(enc.shape[1]):
        jo=m.run_joiner(enc[0,k:k+1], dout); tid=int(jo.argmax())
        if tid!=blank:
            hyp.append(tid); dout=m.run_decoder(hyp[-m.context_size:])
    toks=[id2tok[i] for i in hyp[m.context_size:]]
    return "".join(toks).replace("▁"," ").strip()
def han(s): return "".join(c for c in s if "一"<=c<="鿿")
def _lev(r,h):
    if not r: return 0.0
    d=list(range(len(h)+1))
    for i in range(1,len(r)+1):
        p=d[0]; d[0]=i
        for j in range(1,len(h)+1):
            c=d[j]; d[j]=min(d[j]+1,d[j-1]+1,p+(r[i-1]!=h[j-1])); p=c
    return d[len(h)]/len(r)
def score(ref,hyp):
    is_zh=bool(re.search(r"[一-鿿]",ref))
    if is_zh: return _lev(han(_t2s.convert(_numnorm(ref))),han(_t2s.convert(hyp)))
    return _lev(re.findall(r"[a-z']+",ref.lower()),re.findall(r"[a-z']+",hyp.lower()))