"""Reference-based pronunciation scorer (2-stage RAPS).

Inputs: user audio + reference (standard) audio + canonical text.
Stage 1 (RAPS, end-to-end fine-tuned WavLM + reference cross-attention):
   user/ref audio → V21 forced-align → per-phone user-vs-ref comparison → per-phone engine-style score
Stage 2 (Transformer aggregator):
   per-phone scores → utterance accuracy + integrity (engine-calibrated)

Holdout results (260520, no leakage):
   word   : acc P<10 80.2% / P<15 90.2% ; integ P<10 98.6%
   sentence: acc P<10 77.3% / P<15 89.0% ; integ P<10 82.4% / P<15 93.3%
   all    : acc P<10 78.6% / P<15 89.3% ; integ P<10 87.8% / P<15 94.2%

Models:
   exp_raps_xattn/best.pth      (RAPS stage-1, ~324M)
   exp_stage2_agg/best.pth      (Transformer aggregator, 3.3MB)
   sentence_scorer_v32v3_compact/v21_xlsr_phoneme.pt  (V21 forced align)
   sentence_scorer_v32v3_compact/phone_vocab.json
   gpu_service/custom_phone_dict.json (G2P)

Usage:
   python3 inference_reference.py --user user.wav --ref ref.wav --text "I am going to school"
"""
import os, sys, json, re, argparse
from pathlib import Path
import numpy as np
import torch, torch.nn as nn
import torch.nn.functional as F
import soundfile as sf
from transformers import WavLMModel, Wav2Vec2ForCTC

REPO = Path(__file__).resolve().parent
sys.path.insert(0, str(REPO))
sys.path.insert(0, str(REPO/'sentence_scorer_v32v3_compact'))
SR = 16000; MAX_PHONES = 50

# ---------- G2P ----------
_G2P=None; _CD=None
def g2p():
    global _G2P
    if _G2P is None:
        import g2p_en; _G2P=g2p_en.G2p()
    return _G2P
def custom_dict():
    global _CD
    if _CD is None:
        p=REPO/'gpu_service/custom_phone_dict.json'; _CD=json.load(open(p)) if p.exists() else {}
    return _CD
def text_to_phones(text, phone_to_id):
    """Return phone ids + word_id-per-phone (for fluency word grouping)."""
    text=re.sub(r"[^a-z'\s]"," ",text.lower()); text=re.sub(r"\s+"," ",text).strip()
    ids=[]; wids=[]
    for wi,w in enumerate(text.split()):
        cd=custom_dict()
        phs=cd[w] if w in cd else [re.sub(r"\d+$","",t).lower() for t in g2p()(w) if t!=" "]
        for p in phs:
            ids.append(phone_to_id.get(p, phone_to_id.get('<unk>',1))); wids.append(wi)
    return ids, wids

def load_audio(path, max_sec=10.0):
    w,sr=sf.read(path,dtype='float32')
    if w.ndim>1: w=w.mean(1)
    if sr!=SR:
        idx=np.linspace(0,len(w)-1,int(len(w)*SR/sr)).astype(np.int64); w=w[idx]
    w=w[:int(max_sec*SR)]; w=np.pad(w,int(0.2*SR)).astype(np.float32)
    return (w-w.mean())/(w.std()+1e-7)

# ---------- V21 aligner (reuse) ----------
from inference_sentence import V21Aligner

# ---------- Stage 1: RAPS ----------
class RAPS(nn.Module):
    def __init__(s, n_phon=60, embed=256):
        super().__init__()
        s.w=WavLMModel.from_pretrained('microsoft/wavlm-large'); H=s.w.config.hidden_size
        s.proj=nn.Linear(H*4,embed); s.pe=nn.Embedding(n_phon+1,embed,padding_idx=0)
        s.pos=nn.Parameter(torch.zeros(1,MAX_PHONES,embed))
        s.blocks=nn.ModuleList([nn.TransformerEncoderLayer(embed,8,embed*4,0.1,batch_first=True) for _ in range(4)])
        s.xattn=nn.MultiheadAttention(H,8,batch_first=True)
        s.phn=nn.Linear(embed,1); s.pherr=nn.Linear(embed,1)
        s.utt=nn.Sequential(nn.Linear(embed+7,embed),nn.GELU(),nn.Linear(embed,5))
    def enc(s,w,dev):
        wt=torch.from_numpy(w).unsqueeze(0).to(dev)
        with torch.amp.autocast('cuda',dtype=torch.bfloat16):
            o=s.w(wt,attention_mask=torch.ones(1,len(w),device=dev)).last_hidden_state[0]
        return o.float(), s.w._get_feat_extract_output_lengths(torch.tensor([len(w)])).item()

class AGG(nn.Module):
    """Unified 4-metric aggregator: phone scores + fluency feats → acc/integ/flu/overall."""
    def __init__(s,emb=128,n_phon=60,nflu=8):
        super().__init__(); s.inp=nn.Linear(2,emb); s.pe=nn.Embedding(n_phon+1,emb,padding_idx=0)
        s.pos=nn.Parameter(torch.zeros(1,MAX_PHONES,emb))
        s.bl=nn.ModuleList([nn.TransformerEncoderLayer(emb,8,emb*4,0.1,batch_first=True) for _ in range(4)])
        s.h=nn.Sequential(nn.Linear(emb+nflu,emb),nn.GELU(),nn.Linear(emb,4))
    def forward(s,x,pc,f):
        vm=(pc>0); h=s.inp(x)+s.pe(pc)+s.pos[:,:x.size(1)]
        for b in s.bl: h=b(h,src_key_padding_mask=~vm)
        vmf=vm.float().unsqueeze(-1); pooled=(h*vmf).sum(1)/vmf.sum(1).clamp(min=1)
        return s.h(torch.cat([pooled,f],-1))

def fluency_feats(bu, wids, n):
    """8-dim fluency feats from V21 user spans + word ids — matches unified_aggregator training."""
    s=np.array([bu[i][0] for i in range(n)]); e=np.array([bu[i][1] for i in range(n)]); w=np.array(wids[:n])
    total=e.max()-s.min()+1 if n>0 else 1
    wd=[]; gaps=[]; prev=None
    for wi in sorted(set(w.tolist())):
        ph=np.where(w==wi)[0]; ws=s[ph].min(); we=e[ph].max(); wd.append(we-ws+1)
        if prev is not None: gaps.append(max(0,ws-prev))
        prev=we
    wd=np.array(wd) if wd else np.array([1.]); gaps=np.array(gaps) if gaps else np.array([0.])
    return np.array([n/max(total,1)*50, total/50.0, len(set(w.tolist())),
                     wd.mean()/50.0, wd.std()/50.0, gaps.mean()/50.0, gaps.max()/50.0,
                     float((gaps>10).sum())], np.float32)

def mapb(b,Tv,Tw):
    a=int(b[0]*Tw/max(Tv,1)); bb=max(a,int((b[1]+1)*Tw/max(Tv,1))-1)
    a=max(0,min(a,Tw-1)); bb=max(a,min(bb,Tw-1)); return a,bb

class ReferenceScorer:
    def __init__(s, device='cuda:0',
                 raps='exp_raps_clean/best.pth', agg='exp_unified_agg/best.pth',
                 v21='sentence_scorer_v32v3_compact/v21_xlsr_phoneme.pt',
                 vocab='sentence_scorer_v32v3_compact/phone_vocab.json'):
        s.dev=device; torch.cuda.set_device(int(device.split(':')[1]))
        pv=json.load(open(REPO/vocab)); s.vocab=pv['phone_vocab']; s.p2id=pv['phone_to_id']
        s.aligner=V21Aligner(str(REPO/v21), device)
        s.raps=RAPS().to(device); s.raps.load_state_dict(torch.load(REPO/raps,map_location='cpu',weights_only=False)['model_state'],strict=False); s.raps.eval()
        s.agg=AGG().to(device); s.agg.load_state_dict(torch.load(REPO/agg,map_location='cpu',weights_only=False)['model_state']); s.agg.eval()

    @torch.no_grad()
    def score(s, user_path, ref_path, text):
        ids,wids=text_to_phones(text, s.p2id)
        ids=ids[:MAX_PHONES]; wids=wids[:MAX_PHONES]
        n=len(ids)
        if n==0: return {'error':'empty text'}
        wu=load_audio(user_path); wr=load_audio(ref_path)
        bu,Tu,_=s.aligner.align_and_posteriors(wu,ids)
        br,Tr,_=s.aligner.align_and_posteriors(wr,ids)
        Hu,Lu=s.raps.enc(wu,s.dev); Hr,Lr=s.raps.enc(wr,s.dev)
        toks=[]
        for p in range(n):
            a,b=mapb(bu[p],Tu,Lu); c,d=mapb(br[p],Tr,Lr)
            U=Hu[a:b+1]; R=Hr[c:d+1]
            A,_=s.raps.xattn(U.unsqueeze(0),R.unsqueeze(0),R.unsqueeze(0),need_weights=False); A=A[0]
            u=U.mean(0); am=A.mean(0); mx=(U-A).abs().max(0).values
            toks.append(s.raps.proj(torch.cat([u,am,u-am,mx])))
        x=torch.stack(toks).unsqueeze(0)
        pidt=torch.tensor([[i+1 for i in ids]],device=s.dev)
        x=x+s.raps.pe(pidt)+s.raps.pos[:,:n]
        for blk in s.raps.blocks: x=blk(x)
        phn=(s.raps.phn(x).squeeze(-1)[0]).cpu().numpy()  # 0-10 scale
        phone_scores=(phn*10).clip(0,100)                 # 0-100
        # unified aggregator: phone scores + fluency feats → acc/integ/flu/overall
        sc=np.zeros((1,MAX_PHONES),np.float32); pc=np.zeros((1,MAX_PHONES),np.int64); per=np.zeros((1,MAX_PHONES),np.float32)
        sc[0,:n]=phn; pc[0,:n]=[i+1 for i in ids]
        flu=fluency_feats(bu, wids, n)[None]
        xin=np.stack([sc,per],-1)
        out=s.agg(torch.from_numpy(xin).to(s.dev), torch.from_numpy(pc).to(s.dev),
                  torch.from_numpy(flu).to(s.dev))[0].cpu().numpy()
        clip=lambda v: round(float(np.clip(v,0,100)),1)
        return {
            'accuracy': clip(out[0]),
            'integrity': clip(out[1]),
            'fluency': clip(out[2]),
            'overall': clip(out[3]),
            'phones': [s.vocab[i] for i in ids],
            'phone_scores': [round(float(x),1) for x in phone_scores],
        }

if __name__=='__main__':
    ap=argparse.ArgumentParser()
    ap.add_argument('--user',required=True); ap.add_argument('--ref',required=True)
    ap.add_argument('--text',required=True); ap.add_argument('--gpu',type=int,default=0)
    a=ap.parse_args()
    sc=ReferenceScorer(device=f'cuda:{a.gpu}')
    print(json.dumps(sc.score(a.user,a.ref,a.text),indent=2,ensure_ascii=False))