"""Synthetic full-size frozen-teacher readiness, no retained training model."""
import json,time,hashlib,torch
from train_stable_contextual_rl import ROOT,make_model,make_teacher,tensor_digest
from reference_contextual_operator import accumulate_reference_question
from run_baseline import Qwen3Tokenizer

def main():
 cp=ROOT/'notes/stable_contextual_preflight_manifest.json';c=json.loads(cp.read_text())
 for p,h in c['input_sha256'].items():
  from extract_latent_teacher import sha
  assert sha(ROOT/p)==h,p
 start=time.monotonic();model=make_model(c)
 model.base.load_state_dict(torch.load(ROOT/c['parent'],map_location='cpu',weights_only=True))
 teacher=make_teacher(c);before=tensor_digest(teacher);assert before==tensor_digest(model)
 tok=Qwen3Tokenizer(tokenizer_file_path=ROOT/'code/data/tokenizer-base.json');eos=tok.eos_token_id
 opt=torch.optim.AdamW(model.parameters(),lr=1e-6,weight_decay=.01,foreach=False)
 prompt=torch.randint(3,1000,(960,)).tolist()
 outputs=[torch.randint(3,1000,(n-1,)).tolist()+[eos] for n in [512,192,256,320,384,448,480,512]]
 rewards=[1.,0.,0.,0.,0.,1.,1.,0.];records=[]
 torch.cuda.reset_peak_memory_stats()
 for step,mode in enumerate(['contextual','empty']):
  opt.zero_grad(set_to_none=True)
  r=accumulate_reference_question(model,teacher,tok,prompt,outputs,rewards,[True]*8,context_mode=mode)
  norm=torch.nn.utils.clip_grad_norm_(model.parameters(),1.,error_if_nonfinite=True)
  opt.step();opt.zero_grad(set_to_none=True)
  assert all(p.grad is None for p in teacher.parameters())
  r.update(mode=mode,preclip_norm=float(norm));records.append(r)
  print(json.dumps({'step':step+1,'mode':mode,'seconds':time.monotonic()-start,'norm':float(norm)}),flush=True)
 after=tensor_digest(teacher)
 assert before==after and before!=tensor_digest(model)
 assert all(torch.isfinite(p).all() for p in model.parameters())
 r={'status':'PASS','scope':'Synthetic readiness only; full FP32 student and separate frozen initial Warm teacher, contextual and empty steps with live AdamW buffers; no retained checkpoint or math outputs.','manifest_sha256':hashlib.sha256(cp.read_bytes()).hexdigest(),'teacher_sha256_before':before,'teacher_sha256_after':after,'records':records,'peak_allocated_bytes':torch.cuda.max_memory_allocated(),'peak_reserved_bytes':torch.cuda.max_memory_reserved(),'seconds':time.monotonic()-start}
 (ROOT/'notes/stable_contextual_preflight.json').write_text(json.dumps(r,indent=2));print(json.dumps(r),flush=True)
if __name__=='__main__':main()
