"""Bounded prompt-development experiment on frozen synthetic cases (not held-out)."""
import asyncio,argparse,json,time,sys
from pathlib import Path
import numpy as np
from duplex_cases import CASES
from coreai_d3 import D3CoreAI
from d3_format import render
from d3_runtime import normalize_question

VARIANTS={
 'short_cn':dict(type='choice',instructions='根据对话状态判断如何处理用户刚说的话。注意：对问题的简短回答也是有效输入；附和不需要回复。',criteria={
 'interrupt':'助手还在说话，用户要求停下、暂停、取消，纠正助手，或必须立即改变当前话题。',
 'discard':'噪声、助手回声、对别人说话、单纯附和或感谢。助手正说话时要求继续，或空闲无任务时要求停止，也忽略。',
 'enqueue':'其他有意义的问题、请求、补充或回答；保留并稍后处理。助手空闲时也接纳。',}),
 'short_en':dict(type='choice',instructions='How should a voice assistant handle the latest user utterance? Use the playback state and preceding question. Short answers to a question are meaningful.',criteria={
 'interrupt':'While the assistant is speaking: stop/pause/cancel it, correct what it is saying, urgently switch tasks, or warn of danger.',
 'discard':'Noise, confirmed echo, speech to another person, or acknowledgment needing no reply. Also continue when already speaking, or stop when idle with no task.',
 'enqueue':'A meaningful question, request, information, or answer that can be processed next. When idle, accept meaningful input this way.',}),
}

def state_text(c):
 s=c['state']
 return '助手状态：'+('正在播报' if s['assistant_speaking'] else '当前没有播报')+'\n助手最近说：'+s['last_assistant_utterance']+'\n输入来源：'+{'user':'用户对助手说话','non_speech':'非语音噪声','confirmed_assistant_echo':'已确认是助手回声','not_addressed_to_assistant':'对其他人说话'}[s['input_origin']]+'\n用户本次说：'+s['user_transcript']

def summary(rows):
 labels=['interrupt','discard','enqueue'];conf={x:{y:0 for y in labels} for x in labels}
 for r in rows:conf[r['expected']][r['choice']]+=1
 return dict(cases=len(rows),correct=sum(r['expected']==r['choice'] for r in rows),confusion=conf,false_discard=sum(r['choice']=='discard' and r['expected']!='discard' for r in rows),false_interrupt=sum(r['choice']=='interrupt' and r['expected']!='interrupt' for r in rows),missed_interrupt=sum(r['choice']!='interrupt' and r['expected']=='interrupt' for r in rows),median_seconds=float(np.median([r['seconds'] for r in rows])))

async def main(a):
 engine=await D3CoreAI.load(a.bundle);results={}
 for name,q in VARIANTS.items():
  rows=[]
  for c in CASES:
   s=state_text(c);t=time.perf_counter();logits,keys,encoded=await engine.logits(s,q)
   p=np.exp(logits-logits.max());p/=p.sum()
   r=dict(id=c['id'],expected=c['expected'],choice=keys[int(p.argmax())],probabilities=p.tolist(),seconds=time.perf_counter()-t,tokens=int(len(encoded['input_ids'][0])))
   rows.append(r)
  results[name]=dict(question=q,summary=summary(rows),records=rows)
  Path(a.output).write_text(json.dumps(results,ensure_ascii=False,indent=2)+'\n');print(name,results[name]['summary'],flush=True)

if __name__=='__main__':
 p=argparse.ArgumentParser();p.add_argument('--bundle',default='.');p.add_argument('--output',required=True);a=p.parse_args();asyncio.run(main(a))