"""Complete-record audit and paired native readout; refuses partial endpoints."""
import argparse,collections,json,pathlib,math
import numpy as np
from train_contextual_rl import ROOT,validate
from extract_latent_teacher import sha
from run_baseline import Qwen3Tokenizer,render_prompt,extract_final_candidate
from reasoning_from_scratch.ch03 import normalize_text
from run_with_isolated_grader import IsolatedGrader

def loadlines(p):return [json.loads(s) for s in p.read_text().splitlines()]
def main():
 ap=argparse.ArgumentParser();ap.add_argument('--contract',required=True);a=ap.parse_args();cp=ROOT/a.contract;c=json.loads(cp.read_text());validate(c)
 root=ROOT/c['output'];out=root/'analysis';out.mkdir(exist_ok=False);split=json.loads((ROOT/c['split']).read_text());data=json.loads((ROOT/'code/data/math_train.json').read_text());tok=Qwen3Tokenizer(tokenizer_file_path=ROOT/'code/data/tokenizer-base.json')
 # Existence and structural checks for ALL endpoints precede native extraction.
 raw={};evalwork={}
 for arm in ['W','G','C','E','H']:
  rows=[];work=collections.Counter()
  for j in range(2):
   folder=root/('eval_'+arm);s=json.loads((folder/f'chunk{j}_summary.json').read_text());p=folder/f'chunk{j}_records.jsonl'
   assert s['status']=='COMPLETE' and s['questions']==256 and s['contract_sha256']==sha(cp) and sha(p)==s['records_sha256']
   r=loadlines(p);assert len(r)==256;rows+=r;work.update(s['work'])
  assert len({r['index'] for r in rows})==512 and {r['index'] for r in rows}==set(split['development_indices'])
  assert all(0<len(r['tokens'])<=2048 and tok.eos_token_id not in r['tokens'][:-1] and tok.decode(r['tokens'])==r['text'] for r in rows)
  assert work['generated_tokens']==sum(len(r['tokens']) for r in rows)==work['head_positions']
  raw[arm]=rows;evalwork[arm]=dict(work)
 training={};grader=IsolatedGrader(normalize_text,out/'grading.jsonl')
 try:
  for arm in ['G','C','E','H']:
   allrows=[];allupdates=[];work=collections.Counter();expected_updates=0
   for chunk in range(c['chunks']):
    folder=root/arm;s=json.loads((folder/f'chunk{chunk}_summary.json').read_text());p=folder/f'chunk{chunk}_rollouts.jsonl';u=folder/f'chunk{chunk}_updates.jsonl'
    assert s['status']=='COMPLETE' and s['contract_sha256']==sha(cp) and sha(p)==s['raw_sha256'] and sha(u)==s['updates_sha256']
    allrows+=loadlines(p);allupdates+=loadlines(u)
   assert len(allrows)==256*8 and len(allupdates)==64
   bygroup=collections.defaultdict(list)
   for r in allrows:
    index=split['train_indices'][(r['step']-1)*4+r['position']];assert r['index']==index
    assert r['seed']==c['seed']+(r['step']-1)*4+r['position']
    assert r['prompt']==tok.encode(render_prompt(data[index]['problem'])) and tok.decode(r['tokens'])==r['text']
    assert 0<len(r['tokens'])<=512 and tok.eos_token_id not in r['tokens'][:-1] and r['eos']==(r['tokens'][-1]==tok.eos_token_id)
    pred=extract_final_candidate(r['text'],fallback=None);grade=float(bool(pred) and grader.grade(pred,data[index]['answer']))
    assert pred==r['prediction'] and grade==r['reward'];bygroup[r['step'],r['position']].append(r)
   for step,u in enumerate(allupdates,1):
    assert u['step']==step and len(u['groups'])==4
    calls=0
    for position,g in enumerate(u['groups']):
     rows=sorted(bygroup[step,position],key=lambda r:r['sample']);assert [r['sample'] for r in rows]==list(range(8))
     rr=np.array([r['reward'] for r in rows],dtype=np.float32);eos=np.array([r['eos'] for r in rows]);adv=(rr-rr.mean())/(rr.std(ddof=1)+1e-4)*eos
     assert np.allclose(adv,g['advantages'],atol=3e-7,rtol=3e-7)
     donor=next((i for i,r in enumerate(rows) if r['reward']==1 and r['eos']),None)
     receivers=[] if donor is None else [i for i,r in enumerate(rows) if r['reward']==0 and r['eos']]
     assert donor==g['donor'] and receivers==g['receivers']
     active=[i for i in range(8) if adv[i]!=0 or (arm in ['C','E'] and i in receivers) or (arm=='H' and i==donor)]
     assert len(active)==g['backward_calls'];calls+=len(active)
     bw={'trunk_positions':0,'head_positions':0,'attention_pairs':0}
     for i in active:
      length=len(rows[i]['prompt'])+len(rows[i]['tokens'])-1;bw['trunk_positions']+=length;bw['head_positions']+=len(rows[i]['tokens']);bw['attention_pairs']+=length*length
     for key,v in bw.items():assert g['work'].get('backward_'+key,0)==v
     if arm in ['C','E']:
      assert g['dense_tokens']==sum(len(rows[i]['tokens']) for i in receivers)
     else:assert g['dense_tokens']==0 and g['dense_objective']==0
     for prefix in ['support','teacher']:
      expected={'trunk_positions':0,'head_positions':0,'attention_pairs':0}
      if arm in ['C','E'] and receivers:
       from prefix_teacher_spec import TEMPLATE
       solution=tok.decode(rows[donor]['tokens'][:-1]) if arm=='C' else ''
       tp=tok.encode(TEMPLATE.format(solution=solution,native_prompt=tok.decode(rows[donor]['prompt'])))
       assert len(tp)==g['teacher_prompt_tokens']
       for i in receivers:
        p=rows[i]['prompt'] if prefix=='support' else tp
        length=len(p)+len(rows[i]['tokens'])-1
        expected['trunk_positions']+=length;expected['head_positions']+=len(rows[i]['tokens']);expected['attention_pairs']+=length*length
      for key,v in expected.items():assert g['work'].get(prefix+'_'+key,0)==v
     work.update(g['work'])
    expected_updates+=int(calls>0);assert u['optimizer_updates']==expected_updates and math.isfinite(u['preclip_norm'])
   final=json.loads((root/arm/f"chunk{c['chunks']-1}_summary.json").read_text());assert final['optimizer_updates']==expected_updates
   assert final['work']['generation_generated_tokens']==sum(len(r['tokens']) for r in allrows)
   assert final['work']['questions']==256 and final['work']['correct']==sum(r['reward'] for r in allrows)
   training[arm]={'optimizer_updates':expected_updates,'correct':sum(r['reward'] for r in allrows),'eos':sum(r['eos'] for r in allrows),'work':final['work']}
  values={};answers=[]
  for arm,rows in raw.items():
   byid={r['index']:r for r in rows};bits=[]
   for index in split['development_indices']:
    r=byid[index];pred=extract_final_candidate(r['text'],fallback=None);correct=bool(pred) and grader.grade(pred,data[index]['answer']);bits.append(int(correct));answers.append({'arm':arm,'index':index,'prediction':pred,'correct':bool(correct)})
   values[arm]=np.array(bits)
 finally:grader.close()
 contrasts={}
 for comparator in ['W','G','E','H']:
  delta=values['C']-values[comparator];rng=np.random.default_rng(c['bootstrap_seed']);resamples=np.concatenate([delta[rng.integers(512,size=(1000,512))].mean(1) for _ in range(100)])
  contrasts['C-'+comparator]={'difference':float(delta.mean()),'net_correct':int(delta.sum()),'ci95':np.quantile(resamples,[.025,.975]).tolist(),'gains':int((delta==1).sum()),'losses':int((delta==-1).sum())}
 gate=contrasts['C-W']['net_correct']>=16 and contrasts['C-W']['ci95'][0]>0 and contrasts['C-G']['ci95'][0]>0
 result={'status':'COMPLETE_AUDITED','counts':{k:int(v.sum()) for k,v in values.items()},'contrasts':contrasts,'development_gate':bool(gate),'training':training,'evaluation_work':evalwork,'scope':'Reused512 development; one seed/start; not compute matched; conditional paired item intervals exclude seed uncertainty; no benchmark or sustained-loop claim'}
 (out/'answers.json').write_text(json.dumps(answers));(out/'result.json').write_text(json.dumps(result,indent=2));print(json.dumps(result,indent=2))
if __name__=='__main__':main()
