"""Native greedy endpoint only. Ground-truth answers never enter model inputs."""
import argparse,json,time,collections,torch
from train_contextual_rl import ROOT,validate,make_model
from extract_latent_teacher import sha
from compressed_checkpoint import load_compressed
from run_baseline import Qwen3Tokenizer,render_prompt
from generate_native_policy import generate_policy

def main():
 ap=argparse.ArgumentParser();ap.add_argument('--contract',required=True);ap.add_argument('--arm',choices=['W','G','C','E','H'],required=True);ap.add_argument('--chunk',type=int,choices=[0,1],required=True);a=ap.parse_args()
 cp=ROOT/a.contract;c=json.loads(cp.read_text());validate(c);out=ROOT/c['output']/('eval_'+a.arm);out.mkdir(parents=True,exist_ok=True)
 path=out/f'chunk{a.chunk}_summary.json';assert not path.exists();start=time.monotonic();model=make_model(c).eval()
 if a.arm=='W':
  model.base.load_state_dict(torch.load(ROOT/c['parent'],map_location='cpu',weights_only=True));checkpoint_sha=sha(ROOT/c['parent'])
 else:
  s=json.loads((ROOT/c['output']/a.arm/f"chunk{c['chunks']-1}_summary.json").read_text());assert s['status']=='COMPLETE' and s['completed_steps']==64
  saved=load_compressed(s['checkpoint']['path'],s['checkpoint']);assert saved['contract_sha256']==sha(cp) and saved['arm']==a.arm
  model.base.load_state_dict(saved['model']);del saved;checkpoint_sha=s['checkpoint']['archive_sha256']
 tok=Qwen3Tokenizer(tokenizer_file_path=ROOT/'code/data/tokenizer-base.json');data=json.loads((ROOT/'code/data/math_train.json').read_text());split=json.loads((ROOT/c['split']).read_text())
 entries=sorted([(i,tok.encode(render_prompt(data[i]['problem']))) for i in split['development_indices']],key=lambda x:len(x[1]))[a.chunk*256:(a.chunk+1)*256]
 assert len(entries)==256 and all(len(p)+2048<=4096 for _,p in entries)
 work=collections.Counter();torch.cuda.reset_peak_memory_stats()
 with (out/f'chunk{a.chunk}_records.jsonl').open('x') as log:
  for begin in range(0,256,c['eval_batch']):
   if time.monotonic()-start>3000:raise TimeoutError('Evaluation time bound')
   batch=entries[begin:begin+c['eval_batch']];outputs,w=generate_policy(model,[p for _,p in batch],boundary=None,eos=tok.eos_token_id,pad=tok.eos_token_id,max_tokens=2048)
   work.update(w)
   for (i,p),tokens in zip(batch,outputs):log.write(json.dumps({'index':i,'prompt_tokens':len(p),'tokens':tokens,'text':tok.decode(tokens)})+'\n')
   log.flush();print(json.dumps({'arm':a.arm,'chunk':a.chunk,'completed':begin+len(batch),'seconds':time.monotonic()-start}),flush=True)
 r={'status':'COMPLETE','arm':a.arm,'chunk':a.chunk,'questions':256,'contract_sha256':sha(cp),'checkpoint_sha256':checkpoint_sha,'records_sha256':sha(out/f'chunk{a.chunk}_records.jsonl'),'work':dict(work),'seconds':time.monotonic()-start,'peak_allocated_bytes':torch.cuda.max_memory_allocated(),'scope':'ungraded native512development; no benchmark'}
 path.write_text(json.dumps(r,indent=2));print(json.dumps(r),flush=True)
if __name__=='__main__':main()
