"""TurnSense 1.0 FP32: real causal-prefix inference with the official frontend."""
import os
for k in ('OMP_NUM_THREADS','MKL_NUM_THREADS','OPENBLAS_NUM_THREADS'): os.environ[k]='1'
os.environ['NUMBA_CACHE_DIR']=str(__import__('pathlib').Path(__file__).resolve().parent/'cache/numba')
import argparse,csv,hashlib,json,platform,sys,time
import importlib.metadata as md
from pathlib import Path
from collections import Counter
import numpy as np
import soundfile as sf
import librosa
import onnxruntime as ort
HERE=Path(__file__).resolve().parent
ROOT=HERE.parents[1]
OUT=HERE/'results'
LABEL_PATH=ROOT/'labels/EOT标注/工作数据/eot_labels.csv'
sys.path.insert(0,str(HERE/'vendor'))
from frontend import AudioFrontend
from infer import DEFAULT_FRONTEND_CONF,process_predictions,LABELS

def sha(p): return hashlib.sha256(p.read_bytes()).hexdigest()
def prepare(user,sr,ms):
    end=min(len(user),ms*sr//1000)
    if end<=0: raise ValueError('No audio available')
    prefix=user[:end].astype(np.float32)/32768.
    if sr!=16000: prefix=librosa.resample(prefix,orig_sr=sr,target_sr=16000)
    return np.ascontiguousarray(prefix[-128000:])
class Model:
    def __init__(self):
        manifest=json.loads((HERE/'model-manifest.json').read_text())
        for name,digest in manifest['files'].items():
            if sha(HERE/'models'/name)!=digest: raise ValueError('Weight/CMVN checksum mismatch')
        opts=ort.SessionOptions();opts.intra_op_num_threads=1;opts.inter_op_num_threads=1
        opts.graph_optimization_level=ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        self.session=ort.InferenceSession(str(HERE/'models/model_fp32.onnx'),sess_options=opts,providers=['CPUExecutionProvider'])
        self.frontend=AudioFrontend(**{**DEFAULT_FRONTEND_CONF,'cmvn_file':str(HERE/'models/am.mvn')})
        assert LABELS==['complete','incomplete','invalid']
    def predict(self,audio):
        t=time.perf_counter();feats,n=self.frontend.extract_features(audio)
        inputs={'feats':np.asarray(feats,dtype=np.float32)[None,...],'feat_lengths':np.asarray([int(n)],dtype=np.int64)}
        frontend_ms=(time.perf_counter()-t)*1000
        t=time.perf_counter();raw=self.session.run(None,inputs)[0];inference_ms=(time.perf_counter()-t)*1000
        probs,pred=process_predictions(raw)
        if probs.shape!=(1,3) or not np.isfinite(probs).all() or not np.isclose(probs.sum(),1): raise ValueError('Invalid 3-class output')
        return probs[0],int(pred[0]),frontend_ms,inference_ms

def metrics(rows,rule):
    c=Counter((int(r['can_reply_now']),int(r[rule])) for r in rows if r['can_reply_now'] in ('0','1'))
    tn,fp,fn,tp=c[0,0],c[0,1],c[1,0],c[1,1];n=tn+fp+fn+tp
    return dict(n=n,tn=tn,fp=fp,fn=fn,tp=tp,accuracy=(tn+tp)/n,false_reply_rate=fp/(tn+fp),false_hold_rate=fn/(tp+fn))
def prov():
    return dict(machine=platform.platform(),hardware='Apple M4 Mac mini CPU / 16 GB',python=sys.version,
        dependencies={p:md.version(p) for p in ['numpy','onnxruntime','librosa','kaldi-native-fbank','soundfile']},
        providers=['CPUExecutionProvider'],intra_op_threads=1,inter_op_threads=1,
        model_manifest=json.loads((HERE/'model-manifest.json').read_text()),script_sha256=sha(Path(__file__)),
        annotation_sha256=sha(LABEL_PATH),frontend_sha256=sha(HERE/'vendor/frontend/audio_frontend.py'),
        official_infer_sha256=sha(HERE/'vendor/infer.py'))
def main(mode):
    OUT.mkdir(exist_ok=True)
    t=time.perf_counter();model=Model();load_ms=(time.perf_counter()-t)*1000
    rng=np.random.default_rng(42);test=rng.normal(0,.005,16000).astype(np.float32)
    a=model.predict(test)[0];b=model.predict(test)[0];assert np.array_equal(a,b)
    print('WITNESS CPU seed=42',a.tolist(),'repeat identical',flush=True)
    if mode=='smoke':
        print(json.dumps(dict(load_ms=load_ms,provenance=prov()),ensure_ascii=False));return
    with LABEL_PATH.open(encoding='utf-8-sig',newline='') as f: labels=list(csv.DictReader(f))
    rows,calls=[],[]
    for call in sorted({r['call_id'] for r in labels}):
        path=ROOT/'data/韩泰机器人抢话/音频'/f'{call}.mp3'
        pcm,sr=sf.read(path,dtype='int16',always_2d=True);user=pcm[:,0];duration=round(len(user)*1000/sr,3)
        exact={int(r['decision_ms']):r for r in labels if r['call_id']==call}
        if any(t>duration for t in exact): raise ValueError('Label beyond recording')
        times=sorted(set(range(200,int(duration)+1,200))|set(exact)|{int(duration)}) if mode=='all' else sorted(exact)
        points=[];checkpoints=[]
        for i,ms in enumerate(times):
            t=time.perf_counter();audio=prepare(user,sr,ms);prep=(time.perf_counter()-t)*1000
            probs,pred,front,inf=model.predict(audio);points.append([ms,round(float(probs[0]),7)])
            if ms in exact:
                assert np.array_equal(probs,model.predict(audio)[0])
                modified=user.copy();modified[ms*sr//1000:]=0
                assert np.array_equal(audio,prepare(modified,sr,ms))
                checkpoints.append([ms,round(float(probs[0]),7)])
                rows.append(dict(call_id=call,decision_ms=ms,can_reply_now=exact[ms]['can_reply_now'],
                    p_complete=float(probs[0]),p_incomplete=float(probs[1]),p_invalid=float(probs[2]),
                    native_class=LABELS[pred],native_can_reply=int(pred==0),prediction_05=int(probs[0]>=.5),
                    input_samples=len(audio),available_end_ms=(ms*sr//1000)*1000/sr,
                    preparation_ms=prep,frontend_ms=front,inference_ms=inf))
            if (i+1)%100==0: print(f'{call[:8]}: {i+1}/{len(times)} scores',flush=True)
        calls.append(dict(call_id=call,audio_path=f'韩泰机器人抢话/音频/{call}.mp3',audio_sha256=sha(path),
            duration_ms=duration,points=points,checkpoints=checkpoints))
        print(f'{call[:8]} DONE: {len(points)} scores, {len(checkpoints)} checkpoints; repeat/future-suffix checks passed',flush=True)
    with (OUT/'predictions.csv').open('w',encoding='utf-8',newline='') as f:
        w=csv.DictWriter(f,fieldnames=list(rows[0]));w.writeheader();w.writerows(rows)
    method='TurnSense1.0 FP32, official code/frontend/CMVN and 3-class postprocessing. CPU one intra/inter-op thread. User channel 0 int16/32768, crop received prefix BEFORE librosa resample to 16kHz, last 8s, no future ASR/audio or silence/timeout gates. Score=P(complete), native argmax complete=>1 else0; separately report threshold0.5. Full-prefix resampling is offline causal-window evaluation, not an incremental frontend benchmark.'
    summary=dict(model='TurnSense 1.0 FP32',native_argmax=metrics(rows,'native_can_reply'),threshold_05=metrics(rows,'prediction_05'),
        class_counts=dict(Counter(r['native_class'] for r in rows)),load_ms=load_ms,method=method,provenance=prov(),
        latency_ms={key:{name:float(np.percentile([r[key] for r in rows],q)) for name,q in [('p50',50),('p95',95)]} for key in ['preparation_ms','frontend_ms','inference_ms']})
    (OUT/'summary.json').write_text(json.dumps(summary,ensure_ascii=False,indent=2)+'\n')
    if mode=='all':
        output=dict(schema_version=1,model='turnsense',name='TurnSense 1.0',score='P(complete), 3-class output',sample_interval_ms=200,
            method=method,provenance=prov(),calls=calls)
        (OUT/'score-timeline.json').write_text(json.dumps(output,ensure_ascii=False,separators=(',',':'))+'\n')
    print(json.dumps({k:summary[k] for k in ['native_argmax','threshold_05','class_counts']}),flush=True)
if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('mode',choices=['smoke','evaluate','all']);main(p.parse_args().mode)