"""Fixed four-arm current-policy contextual dense RLVR. No native monitoring."""
import argparse,collections,json,time,pathlib,gc
import torch
from run_baseline import FastModel,QWEN_CONFIG_06_B,Qwen3Tokenizer,render_prompt,extract_final_candidate
from reasoning_from_scratch.ch03 import normalize_text
from latent_deliberation import LatentDeliberation
from contextual_rl_common import accumulate_question,group_advantages,source_roles,teacher_prompt
from contextual_dense_operator import teacher_partition,trajectory_terms,forward_work
from generate_native_policy import generate_policy
from run_with_isolated_grader import IsolatedGrader
from compressed_checkpoint import load_compressed
from broad_latent_common import save_state,rotate_restart
from extract_latent_teacher import sha
ROOT=pathlib.Path(__file__).resolve().parents[1]


def validate(c):
 assert c['status']=='FROZEN_CONTEXTUAL_DENSE_RL'
 for p,h in c['input_sha256'].items():assert sha(ROOT/p)==h,p
 assert c['updates']==64 and c['batch_questions']==4 and c['samples']==8


def make_model(c):
 torch.set_num_threads(4);torch.manual_seed(c['seed'])
 torch.backends.cuda.matmul.allow_tf32=False;torch.backends.cudnn.allow_tf32=False
 base=FastModel(dict(QWEN_CONFIG_06_B,context_length=4096,dtype=torch.float32)).cuda()
 model=LatentDeliberation(base,0,True)
 assert sum(p.numel() for p in model.parameters())==751632384
 return model


@torch.no_grad()
def retention(model,rows,chunk):
 records=[];work=collections.Counter();model.eval()
 for row in rows:
  score,_=trajectory_terms(model,row['prompt'],row['completion'],chunk=chunk)
  records.append({'index':row['index'],'nll':-float(score)/len(row['completion'])})
  work.update(forward_work(row['prompt'],row['completion']))
 return {'rows':records,'mean_nll':sum(r['nll'] for r in records)/len(records),'work':dict(work)}


def component_norms(model,tok,batch,arm,c):
 """Preplanned diagnostics preserve the actual accumulated update exactly."""
 params=list(model.parameters());saved=[None if p.grad is None else p.grad.detach().cpu().clone() for p in params]
 result={};work=collections.Counter()
 try:
  components=['rl']+(['dense'] if arm in ('C','E') else ['replay'] if arm=='H' else [])
  for component in components:
   model.zero_grad(set_to_none=True);calls=0
   for row in batch:
    prompt=row['prompt'];outputs=row['outputs'];adv=group_advantages(row['rewards'],row['finished']);donor,receivers=source_roles(row['rewards'],row['finished'])
    for i,tokens in enumerate(outputs):
     if component=='rl' and adv[i]==0:continue
     if component=='dense' and i not in receivers:continue
     if component=='replay' and i!=donor:continue
     part=None
     if component=='dense':
      tp=teacher_prompt(tok,prompt,outputs[donor],empty=arm=='E')
      part=teacher_partition(model,prompt,tp,tokens,k=c['topk'],chunk=c['head_chunk'])
      for key,v in forward_work(prompt,tokens).items():work['support_'+key]+=v
      for key,v in forward_work(tp,tokens).items():work['teacher_'+key]+=v
     model.train();score,div=trajectory_terms(model,prompt,tokens,part,chunk=c['head_chunk'])
     loss=(-adv[i]*score/(4*8) if component=='rl' else c['dense_weight']*div/(4*8) if component=='dense' else -c['replay_weight']*score/4)
     loss.backward();calls+=1
     for key,v in forward_work(prompt,tokens).items():work['backward_'+key]+=v
   norms=[p.grad.detach().norm() for p in params if p.grad is not None]
   result[component]={'weighted_gradient_norm':float(torch.stack(norms).norm()) if norms else 0.,'backward_calls':calls}
 finally:
  model.zero_grad(set_to_none=True)
  for p,g in zip(params,saved):
   p.grad=None if g is None else g.to(p.device)
   if g is not None:assert torch.equal(p.grad.cpu(),g)
 return {'components':result,'work':dict(work),'actual_gradient_restoration_exact':True}


def main():
 ap=argparse.ArgumentParser();ap.add_argument('--contract',required=True);ap.add_argument('--arm',choices=['G','C','E','H'],required=True);ap.add_argument('--chunk',type=int,required=True);a=ap.parse_args()
 cp=ROOT/a.contract;c=json.loads(cp.read_text());validate(c);assert 0<=a.chunk<c['chunks']
 out=ROOT/c['output']/a.arm;out.mkdir(parents=True,exist_ok=True);summary_path=out/f'chunk{a.chunk}_summary.json';assert not summary_path.exists()
 start=time.monotonic();model=make_model(c);opt=torch.optim.AdamW(model.parameters(),lr=c['learning_rate'],weight_decay=.01,foreach=False)
 work=collections.Counter();previous=None;updates=0
 if a.chunk:
  previous=json.loads((out/f'chunk{a.chunk-1}_summary.json').read_text());assert previous['status']=='COMPLETE'
  saved=load_compressed(previous['checkpoint']['path'],previous['checkpoint']);assert saved['contract_sha256']==sha(cp) and saved['arm']==a.arm
  model.base.load_state_dict(saved['model']);opt.load_state_dict(saved['optimizer']);work.update(saved['work']);updates=saved['optimizer_updates'];del saved
 else:model.base.load_state_dict(torch.load(ROOT/c['parent'],map_location='cpu',weights_only=True))
 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')
 held=[json.loads(line) for line in (ROOT/c['retention_source']).read_text().splitlines()][:64]
 assert not set(split['train_indices'])&set(split['development_indices'])
 if not a.chunk:
  before=retention(model,held,c['head_chunk']);(out/'retention_before.json').write_text(json.dumps(before,indent=2))
  for k,v in before['work'].items():work['retention_'+k]+=v
 grader=IsolatedGrader(normalize_text,out/f'chunk{a.chunk}_grader.jsonl');torch.cuda.reset_peak_memory_stats()
 try:
  with (out/f'chunk{a.chunk}_rollouts.jsonl').open('x') as raw,(out/f'chunk{a.chunk}_updates.jsonl').open('x') as log:
   for step in range(a.chunk*c['updates_per_chunk'],(a.chunk+1)*c['updates_per_chunk']):
    if time.monotonic()-start>2800:raise TimeoutError('Chunk exceeds timing bound; do not silently retry')
    opt.zero_grad(set_to_none=True);batch=[];group_records=[]
    for position in range(4):
     offset=4*step+position;index=split['train_indices'][offset];prompt=tok.encode(render_prompt(data[index]['problem']));seed=c['seed']+offset
     assert len(prompt)<=c['prompt_limit'];torch.manual_seed(seed)
     outputs,gw=generate_policy(model,[prompt]*8,boundary=None,eos=tok.eos_token_id,pad=tok.eos_token_id,max_tokens=c['train_tokens'],temperature=.8,top_p=.9)
     rewards=[];finished=[]
     for sample,tokens in enumerate(outputs):
      text=tok.decode(tokens);prediction=extract_final_candidate(text,fallback=None);reward=float(bool(prediction) and grader.grade(prediction,data[index]['answer']));eos=tokens[-1]==tok.eos_token_id
      rewards.append(reward);finished.append(eos)
      raw.write(json.dumps({'index':index,'step':step+1,'position':position,'sample':sample,'seed':seed,'prompt':prompt,'tokens':tokens,'text':text,'prediction':prediction,'reward':reward,'eos':eos})+'\n')
     batch.append({'index':index,'prompt':prompt,'outputs':outputs,'rewards':rewards,'finished':finished});raw.flush()
     for k,v in gw.items():work['generation_'+k]+=v
     work['questions']+=1;work['correct']+=sum(rewards);work['eos']+=sum(finished)
    for row in batch:
     r=accumulate_question(model,tok,row['prompt'],row['outputs'],row['rewards'],row['finished'],arm=a.arm,batch_questions=4,dense_weight=c['dense_weight'],replay_weight=c['replay_weight'],topk=c['topk'],chunk=c['head_chunk'])
     r['index']=row['index'];group_records.append(r);work.update(r['work'])
    diagnostic=None
    if step+1 in c['gradient_diagnostic_updates']:
     diagnostic=component_norms(model,tok,batch,a.arm,c)
     for k,v in diagnostic['work'].items():work['diagnostic_'+k]+=v
    calls=sum(r['backward_calls'] for r in group_records);norm=0.
    if calls:
     norm=float(torch.nn.utils.clip_grad_norm_(model.parameters(),1.,error_if_nonfinite=True));opt.step();updates+=1
    opt.zero_grad(set_to_none=True)
    record={'step':step+1,'groups':group_records,'diagnostic':diagnostic,'preclip_norm':norm,'optimizer_updates':updates,'seconds':time.monotonic()-start}
    log.write(json.dumps(record)+'\n');log.flush();print(json.dumps({'arm':a.arm,'step':step+1,'correct':work['correct'],'eos':work['eos'],'preclip_norm':norm,'seconds':time.monotonic()-start}),flush=True)
 finally:grader.close()
 final=a.chunk==c['chunks']-1
 if final:
  after=retention(model,held,c['head_chunk']);(out/'retention_after.json').write_text(json.dumps(after,indent=2))
  for k,v in after['work'].items():work['retention_'+k]+=v
 assert all(torch.isfinite(p).all() for p in model.parameters())
 payload={'model':{k:v.detach().cpu() for k,v in model.base.state_dict().items()},'arm':a.arm,'contract_sha256':sha(cp),'completed_steps':(a.chunk+1)*c['updates_per_chunk'],'optimizer_updates':updates,'work':dict(work)}
 if not final:payload['optimizer']=opt.state_dict()
 checkpoint=save_state(payload,out/('final.pt.zst' if final else f'restart_{a.chunk}.pt.zst'),final=final)
 result={'status':'COMPLETE','arm':a.arm,'chunk':a.chunk,'contract_sha256':sha(cp),'completed_steps':payload['completed_steps'],'checkpoint':checkpoint,'work':dict(work),'optimizer_updates':updates,'seconds_including_setup_save':time.monotonic()-start,'peak_allocated_bytes':torch.cuda.max_memory_allocated(),'raw_sha256':sha(out/f'chunk{a.chunk}_rollouts.jsonl'),'updates_sha256':sha(out/f'chunk{a.chunk}_updates.jsonl')}
 summary_path.write_text(json.dumps(result,indent=2));rotate_restart(out,previous,result);print(json.dumps(result),flush=True)
if __name__=='__main__':main()
