"""Reproducible native/reference comparison. Only synthetic bundled cases are used."""
import argparse,asyncio,json,platform,resource,time,sys
from pathlib import Path
import numpy as np
from validation_cases import CASES

def write(path,value):
    path.write_text(json.dumps(value,ensure_ascii=False,indent=2,allow_nan=False)+'\n')

def softmax(x):
    x=np.asarray(x,dtype=np.float64); z=np.exp(x-x.max());return z/z.sum()

async def native(a):
    from coreai_d3 import D3CoreAI
    root=Path(a.bundle); evidence=Path(a.evidence); evidence.mkdir(exist_ok=True,parents=True)
    t=time.perf_counter();engine=await D3CoreAI.load(root); load_s=time.perf_counter()-t
    records=[]
    for c in CASES:
        t=time.perf_counter();logits,keys,encoded=await engine.logits(c['state'],c['question']);elapsed=time.perf_counter()-t
        probs=softmax(logits/engine.config['temperature'])
        assert np.isfinite(probs).all() and abs(probs.sum()-1)<1e-6
        arrays={k:np.asarray(v) for k,v in encoded.items() if k in ['input_ids','attention_mask','pixel_values','pixel_values_videos','image_grid_thw','video_grid_thw','mm_token_type_ids']}
        np.savez(evidence/(c['id']+'.npz'),**arrays)
        r=dict(id=c['id'],modality='text',type=c['question']['type'],keys=keys,tokens=int(len(encoded['input_ids'][0])),probabilities=probs.tolist(),choice=keys[int(probs.argmax())],end_to_end_seconds=elapsed)
        r.update(expected=c.get('expected'),suite=c.get('suite','general'))
        records.append(r);write(evidence/'native.json',records); print(json.dumps(r,ensure_ascii=False),flush=True)
    checks={}
    q=CASES[0]['question']
    for label,fn in [
        ('oversized_text_rejected',lambda:engine.logits('refund '*2000,q)),
        ('media_rejected',lambda:engine.logits('x',q,images=['unsupported.png'])),
        ('empty_options_rejected',lambda:engine.logits('x',dict(type='choice',criteria={}))),
        ('unknown_type_rejected',lambda:engine.logits('x',dict(type='unknown'))),
        ('invalid_score_rejected',lambda:engine.logits('x',dict(type='score',criteria=['one']))),
    ]:
        try:await fn()
        except ValueError:checks[label]=True
        else:raise AssertionError(label)
    actual=await engine.system_one(state='The customer has a receipt.',questions={'yes':dict(type='noul',instructions='Does the customer have a receipt?'),'route':q})
    checks['typed_multi_question']=set(actual['answers'])=={'yes','route'}
    repeated=[];times=[]
    for _ in range(5):
        t=time.perf_counter();logits,_,_=await engine.logits(CASES[0]['state'],q);times.append(time.perf_counter()-t);repeated.append(softmax(logits))
    drift=max(float(abs(x-repeated[0]).max()) for x in repeated)
    checks['repeat_probability_max_error']=drift
    assert drift<1e-5
    assert all(v for k,v in checks.items() if k!='repeat_probability_max_error')
    write(evidence/'runtime.json',dict(platform=platform.platform(),load_seconds=load_s,warm_end_to_end_seconds=times,warm_median_seconds=float(np.median(times)),peak_rss_bytes=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss,checks=checks,typed_example=actual))

def reference(a):
    import torch
    sys.path.insert(0,str(Path(a.source).resolve()))
    from d3_runtime import D3
    torch.set_num_threads(4);evidence=Path(a.evidence)
    ref=D3.from_pretrained(a.source,device='cpu',verify='full')
    native_records=json.loads((evidence/'native.json').read_text());records=[]
    for r in native_records:
        e=dict(np.load(evidence/(r['id']+'.npz')));count=len(r['keys']);t=time.perf_counter()
        with torch.inference_mode():
            expected=np.array(ref.probabilities([e['input_ids'][0].tolist()],[count])[0])
        actual=np.array(r['probabilities']);out=dict(r,reference_probabilities=expected.tolist(),reference_choice=r['keys'][int(expected.argmax())],max_probability_error=float(abs(actual-expected).max()),reference_seconds=time.perf_counter()-t)
        out['choice_matches']=out['choice']==out['reference_choice'];records.append(out)
        write(evidence/'comparison.json',records);print(json.dumps(out,ensure_ascii=False),flush=True)
    summary=dict(cases=len(records),choices_matched=sum(r['choice_matches'] for r in records),max_probability_error=max(r['max_probability_error'] for r in records),reference_peak_rss_bytes=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss,limitations=['Synthetic conversion-fidelity tests, not a quality or calibration benchmark.','Text-only artifact; images and video are rejected.','No iPhone or accelerator-placement verification.'])
    duplex=[r for r in records if r['expected'] is not None]
    labels=['interrupt','discard','enqueue'];confusion={x:{y:0 for y in labels} for x in labels}
    for r in duplex:confusion[r['expected']][r['choice']]+=1
    summary['business']=dict(cases=len(duplex),correct=sum(r['choice']==r['expected'] for r in duplex),reference_correct=sum(r['reference_choice']==r['expected'] for r in duplex),confusion=confusion,false_discard=sum(r['choice']=='discard' and r['expected']!='discard' for r in duplex),false_interrupt=sum(r['choice']=='interrupt' and r['expected']!='interrupt' for r in duplex),missed_interrupt=sum(r['choice']!='interrupt' and r['expected']=='interrupt' for r in duplex))
    write(evidence/'summary.json',summary);print(json.dumps(summary),flush=True)
    if summary['choices_matched']!=summary['cases']:sys.exit(2)

if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('phase',choices=['native','reference']);p.add_argument('--bundle',default='.');p.add_argument('--source',required=True);p.add_argument('--evidence',required=True);a=p.parse_args()
    asyncio.run(native(a)) if a.phase=='native' else reference(a)