"""Frozen initial-teacher contextual RLVR with longer ordinary control. No native monitoring."""
import argparse,collections,json,time,pathlib,gc
import torch
import hashlib
from reference_contextual_operator import accumulate_reference_question
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_STABLE_CONTEXTUAL_RL'
 for p,h in c['input_sha256'].items():assert sha(ROOT/p)==h,p
 assert c['updates']=={'G':80,'F':64,'E':64} and c['batch_questions']==4 and c['samples']==8
 assert c['chunks']=={'G':5,'F':4,'E':4} and c['updates_per_chunk']==16


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 tensor_digest(model):
 h=hashlib.sha256()
 for name,value in model.base.state_dict().items():
  h.update(name.encode());h.update(str(value.dtype).encode());h.update(str(tuple(value.shape)).encode())
  h.update(value.detach().cpu().contiguous().numpy().tobytes())
 return h.hexdigest()


def make_teacher(c):
 teacher=make_model(c).eval()
 teacher.base.load_state_dict(torch.load(ROOT/c['parent'],map_location='cpu',weights_only=True))
 teacher.requires_grad_(False)
 assert not any(p.requires_grad for p in teacher.parameters())
 return teacher


def main():
 ap=argparse.ArgumentParser();ap.add_argument('--contract',required=True);ap.add_argument('--arm',choices=['G','F','E'],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'][a.arm]
 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=c['weight_decay'],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))
 teacher=make_teacher(c) if a.arm in ('F','E') else None
 teacher_before=tensor_digest(teacher) if teacher is not None else None
 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='G',batch_questions=4,chunk=c['head_chunk']) if a.arm=='G' else accumulate_reference_question(model,teacher,tok,row['prompt'],row['outputs'],row['rewards'],row['finished'],context_mode='contextual' if a.arm=='F' else 'empty',batch_questions=4,dense_weight=c['dense_weight'],topk=c['topk'],chunk=c['head_chunk']))
     r['index']=row['index'];group_records.append(r);work.update(r['work'])
    calls=sum(r['backward_calls'] for r in group_records);norm=0.
    if calls:
     norm=float(torch.nn.utils.clip_grad_norm_(model.parameters(),c['clip_norm'],error_if_nonfinite=True));opt.step();updates+=1
    opt.zero_grad(set_to_none=True)
    record={'step':step+1,'groups':group_records,'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'][a.arm]-1
 teacher_after=tensor_digest(teacher) if teacher is not None else None
 assert teacher_before==teacher_after
 if teacher is not None:
  assert all(p.grad is None and not p.requires_grad for p in teacher.parameters())
  del teacher;gc.collect();torch.cuda.empty_cache()
 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(),'teacher_tensor_sha256_before':teacher_before,'teacher_tensor_sha256_after':teacher_after,'teacher_exact_unchanged':teacher_before==teacher_after,'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()
