"""Fit-only, SFT-anchored preference training on an OCR model's own mistakes.

No reference model, external text, validation labels, or manual label corrections.
The run has one fixed final checkpoint. All heavy state remains on the GPU host.
"""
from __future__ import annotations
import argparse
import csv
import hashlib
import json
import math
import os
from pathlib import Path
import random
import time

ROOT=Path('/marimo/storage/road')
RUN=ROOT/'runs/visual_preference_fit_oct2'

def digest(path):
 h=hashlib.sha256()
 with Path(path).open('rb') as f:
  for b in iter(lambda:f.read(8*1024*1024),b''):h.update(b)
 return h.hexdigest()

def write_json(path,value):
 p=Path(path);p.parent.mkdir(parents=True,exist_ok=True)
 tmp=p.with_suffix(p.suffix+'.tmp');tmp.write_text(json.dumps(value,indent=2,sort_keys=True,allow_nan=False)+'\n');tmp.replace(p)

def sequence_logps(logits,labels):
 """Completion-only mean log probabilities including EOS; avoid prefix softmax."""
 import torch
 import torch.nn.functional as F
 target=labels[:,1:];mask=target!=-100
 assert mask.any(1).all(),'Missing assistant completion mask'
 losses=F.cross_entropy(logits[:,:-1][mask].float(),target[mask],reduction='none')
 row=torch.arange(target.shape[0],device=target.device)[:,None].expand_as(target)[mask]
 sums=torch.zeros(target.shape[0],device=logits.device).scatter_add(0,row,-losses)
 return sums/mask.sum(1)

def anchored_loss(chosen,rejected,eligible,beta=2.,margin=1.,weight=.1):
 import torch.nn.functional as F
 sft=-chosen.mean()
 preference=-F.logsigmoid(beta*(chosen-rejected)-margin)
 pref=(preference*eligible).sum()/eligible.sum().clamp_min(1)
 return sft+weight*pref, sft.detach(),pref.detach()

def policy():
 freeze=json.loads((ROOT/'visual_preference_oct2.freeze.json').read_text())
 for name,sha in freeze['sources_sha256'].items():assert digest(ROOT/name)==sha,name
 bundle=ROOT/'medusa_fv_gate_bundle'
 for name,sha in freeze['inputs_sha256'].items():assert digest(bundle/name)==sha,name
 split=json.loads((ROOT/'split_manifest.json').read_text())
 assert len(split['fit_ids'])==3577 and len(split['val_ids'])==500
 assert not set(split['fit_ids'])&set(split['val_ids'])
 rows=list(csv.DictReader((ROOT/'Train.csv').open(encoding='utf-8-sig')))
 source={r['ID']:r['Target'] for r in rows};assert len(source)==len(rows)
 assert not {source[i] for i in split['fit_ids']}&{source[i] for i in split['val_ids']}
 sample=random.Random(freeze['configuration']['seed']).sample(split['fit_ids'],freeze['configuration']['fit_sample'])
 return freeze,split,source,sample

def load(freeze,adapter,trainable):
 import torch
 from transformers import AutoModelForImageTextToText,AutoProcessor
 from peft import PeftModel
 from peft.utils import get_peft_model_state_dict
 from safetensors.torch import load_file
 cfg=freeze['configuration']
 processor=AutoProcessor.from_pretrained(freeze['model'],revision=freeze['revision'],min_pixels=64*28*28,max_pixels=cfg['max_pixels'])
 model=AutoModelForImageTextToText.from_pretrained(freeze['model'],revision=freeze['revision'],dtype=torch.bfloat16).cuda()
 visual=[m for n,m in model.named_modules() if n.split('.')[-1]=='visual'];assert len(visual)==1
 visual[0].float() # Establish FP32 master dtype BEFORE the checkpoint is loaded.
 model=PeftModel.from_pretrained(model,str(adapter),is_trainable=trainable)
 saved=load_file(str(Path(adapter)/'adapter_model.safetensors'),device='cpu')
 exported=get_peft_model_state_dict(model)
 assert set(saved)==set(exported),'Incomplete adapter state'
 for name,value in saved.items():
  assert torch.equal(value,exported[name].detach().cpu().to(value.dtype)),name
 del saved,exported
 for name,p in model.named_parameters():
  p.requires_grad_(trainable and 'lora_' in name and '.default.' in name and '.visual.' not in name)
 assert not any(p.requires_grad for n,p in model.named_parameters() if '.visual.' in n)
 if trainable:
  assert any(p.requires_grad for p in model.parameters())
  model.config.use_cache=False
  model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant':False})
  model.enable_input_require_grads();model.train()
 else:model.eval()
 return processor,model

def rng_state():
 import torch,numpy as np
 return {'python':random.getstate(),'numpy':np.random.get_state(),'torch':torch.get_rng_state(),'cuda':torch.cuda.get_rng_state_all()}

def restore_rng(state):
 import torch,numpy as np
 random.setstate(state['python']);np.random.set_state(state['numpy']);torch.set_rng_state(state['torch']);torch.cuda.set_rng_state_all(state['cuda'])

def save_checkpoint(model,processor,optimizer,scheduler,step,manifest,final=False):
 import torch
 name='final' if final else f'checkpoint-{step}'
 destination=RUN/name;temporary=RUN/('.'+name+'.partial')
 assert not destination.exists() and not temporary.exists(),'Refusing to overwrite checkpoint'
 temporary.mkdir(parents=True)
 model.save_pretrained(temporary);processor.save_pretrained(temporary)
 torch.save({'optimizer':optimizer.state_dict(),'scheduler':scheduler.state_dict(),'rng':rng_state(),'step':step},temporary/'training_state.pt')
 write_json(temporary/'run_config.json',manifest)
 write_json(temporary/'SAVE_COMPLETE',{'step':step,'weight_sha256':digest(temporary/'adapter_model.safetensors'),
              'optimizer_state_sha256':digest(temporary/'training_state.pt'),'manifest_sha256':digest(temporary/'run_config.json')})
 temporary.replace(destination)

def make_batch(processor,rows,source,negatives):
 import road_lib as R
 from train_lora import Collator
 images=[R.load_line_image(str(ROOT/'images'/(i+'.jpg')),128,False) for i in rows]
 examples=[{'image':im,'target':source[i]} for i,im in zip(rows,images)]
 examples.extend({'image':im,'target':negatives[i]} for i,im in zip(rows,images))
 processor.tokenizer.padding_side='right'
 batch=Collator(processor)(examples)
 for image in images:image.close()
 return {k:v.cuda() if hasattr(v,'cuda') else v for k,v in batch.items()}

def main():
 ap=argparse.ArgumentParser(description=__doc__);ap.add_argument('--mode',choices=['preflight','mine','train','val_initial','val_final'],required=True)
 a=ap.parse_args();freeze,split,source,ids=policy();cfg=freeze['configuration']
 import torch,numpy as np
 from transformers import get_cosine_schedule_with_warmup
 import road_lib as R
 torch.set_num_threads(4);random.seed(cfg['seed']);np.random.seed(cfg['seed']);torch.manual_seed(cfg['seed']);torch.cuda.manual_seed_all(cfg['seed'])
 assert digest(ROOT/'preference_initial_oct2/adapter_model.safetensors')==freeze['initialization']['weight_sha256']
 adapter=RUN/'final' if a.mode=='val_final' else ROOT/'preference_initial_oct2'
 checkpoint=None
 if a.mode=='train':
  checkpoints=sorted(RUN.glob('checkpoint-*'),key=lambda p:int(p.name.split('-')[-1]))
  checkpoints=[p for p in checkpoints if (p/'SAVE_COMPLETE').exists()]
  if (RUN/'final/SAVE_COMPLETE').exists():print('ALREADY_COMPLETE');return
  if checkpoints:checkpoint=checkpoints[-1];adapter=checkpoint
 processor,model=load(freeze,adapter,a.mode in ['train','preflight'])
 if a.mode in ['mine','val_initial','val_final']:
  processor.tokenizer.padding_side='left'
  chat=processor.apply_chat_template([{'role':'user','content':[{'type':'image'},{'type':'text','text':__import__('train_lora').PROMPT}]}],tokenize=False,add_generation_prompt=True)
  validation=a.mode!='mine';use_ids=split['val_ids'] if validation else ids
  output=ROOT/(f'preds_val_visual_preference_{a.mode[4:]}_oct2.json' if validation else 'visual_preference_negatives_oct2.json')
  identity={'mode':a.mode,'source_sha256':digest(__file__),'adapter_sha256':digest(adapter/'adapter_model.safetensors'),
      'freeze_sha256':digest(ROOT/'visual_preference_oct2.freeze.json'),'ids':use_ids,
      'image_sha256':{i:digest(ROOT/'images'/(i+'.jpg')) for i in use_ids}}
  sidecar=output.with_suffix('.manifest.json')
  if sidecar.exists():assert json.loads(sidecar.read_text())==identity
  else:assert not output.exists();write_json(sidecar,identity)
  records=json.loads(output.read_text()) if output.exists() else {}
  assert set(records)<=set(use_ids)
  todo=[i for i in use_ids if i not in records];bs=4 if not validation else 8
  for start in range(0,len(todo),bs):
   batch=todo[start:start+bs];images=[R.load_line_image(str(ROOT/'images'/(i+'.jpg')),cfg['height'],False) for i in batch]
   inputs=processor(text=[chat]*len(images),images=images,padding=True,return_tensors='pt').to('cuda')
   # Per-batch independent seed makes interrupted mining reproducible.
   torch.manual_seed(cfg['seed']+use_ids.index(batch[0]));torch.cuda.manual_seed_all(cfg['seed']+use_ids.index(batch[0]))
   with torch.inference_mode(),torch.autocast('cuda',dtype=torch.bfloat16):
    generated=model.generate(**inputs,max_new_tokens=cfg['max_new_tokens'],do_sample=not validation,
       **({'temperature':.8,'top_p':.95} if not validation else {'num_beams':cfg['decode_beams']}))
   texts=processor.batch_decode(generated[:,inputs['input_ids'].shape[1]:],skip_special_tokens=True)
   for i,text in zip(batch,texts):records[i]=R.clean_prediction(text,strip_preamble=False)
   for image in images:image.close()
   write_json(output,records)
   write_json(output.with_suffix('.progress.json'),{'completed':len(records),'total':len(use_ids)})
   if start%(bs*10)==0:print('DECODE',a.mode,len(records),flush=True)
  if validation:
   with output.with_suffix('.csv').open('w',newline='') as f:
    w=csv.writer(f);w.writerow(['ID','raw']);w.writerows((i,records[i]) for i in use_ids)
  print('DECODE_COMPLETE',a.mode,flush=True);return
 trainable=[p for p in model.parameters() if p.requires_grad]
 optimizer=torch.optim.AdamW(trainable,lr=cfg['language_lr'],weight_decay=cfg['weight_decay'])
 scheduler=get_cosine_schedule_with_warmup(optimizer,cfg['warmup_steps'],cfg['steps'])
 if a.mode=='preflight':
  checked=ids[:2];negatives={i:('x' if source[i][0]!='x' else 'y')+source[i][1:] for i in checked}
  batch=make_batch(processor,checked,source,negatives);labels=batch.pop('labels')
  visual_before={n:p.detach().cpu().clone() for n,p in model.named_parameters() if '.visual.' in n and '.modules_to_save.' in n}
  with torch.autocast('cuda',dtype=torch.bfloat16):
   logp=sequence_logps(model(**batch).logits,labels);loss,_,_=anchored_loss(logp[:2],logp[2:],torch.ones(2,device='cuda'))
  loss.backward();assert torch.isfinite(loss)
  grads=[p.grad for p in trainable if p.grad is not None]
  assert grads and all(torch.isfinite(g).all() for g in grads)
  assert any(g.abs().sum()>0 for g in grads)
  assert all(p.grad is None for n,p in model.named_parameters() if '.visual.' in n)
  optimizer.step();scheduler.step()
  assert all(torch.equal(p.detach().cpu(),visual_before[n]) for n,p in model.named_parameters() if n in visual_before)
  # An actual save/load checks adapter export and optimizer state before production updates.
  canary=ROOT/'visual_preference_canary_oct2';assert not canary.exists();canary.mkdir()
  model.save_pretrained(canary)
  from safetensors.torch import load_file
  from peft.utils import get_peft_model_state_dict,set_peft_model_state_dict
  state=load_file(str(canary/'adapter_model.safetensors'))
  exported=get_peft_model_state_dict(model)
  assert set(state)==set(exported) and all(torch.equal(state[n],exported[n].detach().cpu()) for n in state)
  optfile=canary/'training_state.pt';torch.save({'optimizer':optimizer.state_dict(),'scheduler':scheduler.state_dict(),'rng':rng_state()},optfile)
  resumed=torch.optim.AdamW(trainable,lr=cfg['language_lr']);resumed.load_state_dict(torch.load(optfile,weights_only=False)['optimizer'])
  assert all(torch.equal(optimizer.state[p]['exp_avg'],resumed.state[p]['exp_avg']) for p in trainable if p in optimizer.state)
  write_json(ROOT/'visual_preference_preflight_oct2.json',{'passed':True,'fit_images':2,'validation_used':False,
   'finite_nonzero_language_gradients':True,'vision_frozen_exactly':True,'adapter_export_exact':True,
   'optimizer_roundtrip_exact':True,'full_resume_next_update_tested':False,
   'gpu_peak_bytes':torch.cuda.max_memory_allocated(),'source_sha256':digest(__file__)})
  print('PREFERENCE_PREFLIGHT_PASSED',flush=True);return
 assert json.loads((ROOT/'visual_preference_preflight_oct2.json').read_text())['source_sha256']==digest(__file__)
 negatives=json.loads((ROOT/'visual_preference_negatives_oct2.json').read_text());assert set(negatives)==set(ids)
 from rapidfuzz.distance import Levenshtein
 eligible={i:cfg['minimum_negative_character_edits']<=Levenshtein.distance(source[i],negatives[i])<=cfg['maximum_negative_character_edits'] for i in ids}
 manifest={'model':freeze['model'],'revision':freeze['revision'],'configuration':cfg,
  'initialization':freeze['initialization'],'training_source_sha256':digest(__file__),
  'freeze_sha256':digest(ROOT/'visual_preference_oct2.freeze.json'),'fit_ids':ids,'val_ids':split['val_ids'],
  'fit_image_sha256':{i:digest(ROOT/'images'/(i+'.jpg')) for i in ids},'validation_contaminated':False,
  'source_csv_sha256':digest(ROOT/'Train.csv'),'split_manifest_sha256':digest(ROOT/'split_manifest.json'),
  'negatives_sha256':digest(ROOT/'visual_preference_negatives_oct2.json'),'eligible_preference_pairs':sum(eligible.values()),
  'trainable_names':[n for n,p in model.named_parameters() if p.requires_grad]}
 RUN.mkdir(parents=True,exist_ok=True)
 if (RUN/'run_config.json').exists():assert json.loads((RUN/'run_config.json').read_text())==manifest
 else:write_json(RUN/'run_config.json',manifest)
 step=0
 if checkpoint:
  marker=json.loads((checkpoint/'SAVE_COMPLETE').read_text());assert marker['weight_sha256']==digest(checkpoint/'adapter_model.safetensors')
  assert marker['optimizer_state_sha256']==digest(checkpoint/'training_state.pt')
  assert json.loads((checkpoint/'run_config.json').read_text())==manifest
  state=torch.load(checkpoint/'training_state.pt',map_location='cpu',weights_only=False)
  optimizer.load_state_dict(state['optimizer']);scheduler.load_state_dict(state['scheduler']);restore_rng(state['rng']);step=state['step']
 order=ids.copy();random.Random(cfg['seed']+1).shuffle(order)
 assert len(order)==cfg['steps']*cfg['batch_pairs']*cfg['gradient_accumulation']
 for current in range(step,cfg['steps']):
  optimizer.zero_grad(set_to_none=True);metrics=[]
  for micro in range(cfg['gradient_accumulation']):
   start=(current*cfg['gradient_accumulation']+micro)*cfg['batch_pairs'];batch_ids=order[start:start+cfg['batch_pairs']]
   batch=make_batch(processor,batch_ids,source,negatives);labels=batch.pop('labels');count=len(batch_ids)
   with torch.autocast('cuda',dtype=torch.bfloat16):
    lp=sequence_logps(model(**batch).logits,labels)
    active=torch.tensor([eligible[i] for i in batch_ids],device='cuda',dtype=torch.float32)
    loss,sft,pref=anchored_loss(lp[:count],lp[count:],active,cfg['preference_beta'],cfg['preference_margin'],cfg['preference_weight'])
   assert torch.isfinite(loss);(loss/cfg['gradient_accumulation']).backward()
   metrics.append((float(sft),float(pref)))
  norm=torch.nn.utils.clip_grad_norm_(trainable,1.);assert torch.isfinite(norm)
  optimizer.step();scheduler.step();completed=current+1
  record={'step':completed,'sft_loss':sum(m[0] for m in metrics)/len(metrics),'preference_loss':sum(m[1] for m in metrics)/len(metrics),'gradient_norm':float(norm),'lr':scheduler.get_last_lr()[0]}
  with (RUN/'loss.jsonl').open('a') as f:f.write(json.dumps(record)+'\n')
  print(json.dumps(record),flush=True)
  if completed%cfg['save_every']==0:save_checkpoint(model,processor,optimizer,scheduler,completed,manifest)
 save_checkpoint(model,processor,optimizer,scheduler,cfg['steps'],manifest,final=True)
 print('PREFERENCE_TRAIN_COMPLETE',flush=True)

if __name__=='__main__':main()