import copy,json,time
from pathlib import Path
import torch
from safetensors.torch import save_file
from structured_policy import StructuredPolicy,FEATURES,encode

torch.manual_seed(42); torch.set_num_threads(4)
device='cuda';out=Path('release/structured');out.mkdir(parents=True,exist_ok=True)
def load(split,n):
    rows=[json.loads(x) for x in open('data/sft/'+split+'.jsonl')][:n]
    x=torch.zeros(len(rows),4,15);mask=torch.zeros(len(rows),4,dtype=torch.bool);target=torch.zeros(len(rows),4)
    for i,r in enumerate(rows):
        u=json.loads(r['messages'][1]['content'].split('\n\nRequested field:')[0]);state=json.loads(u['context'])
        k=len(r['keys']);x[i,:k]=torch.tensor(encode(state,r['keys']));mask[i,:k]=True;target[i,:k]=torch.tensor(r['target'])
    return x.to(device),mask.to(device),target.to(device)
x,mask,target=load('train',4096);vx,vm,vt=load('val',512)
model=StructuredPolicy().to(device)
opt=torch.optim.AdamW(model.parameters(),lr=.001,weight_decay=.0001)
best=float('inf');history=[];start=time.monotonic()
for epoch in range(1,101):
    model.train();order=torch.randperm(len(x),device=device)
    for i in range(0,len(x),128):
        j=order[i:i+128];logp=model(x[j],mask[j]).log_softmax(-1)
        loss=-(target[j]*logp).sum(-1).mean();loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(),1.);opt.step();opt.zero_grad(set_to_none=True)
    if epoch%10==0:
        model.eval()
        with torch.no_grad():
            logits=model(vx,vm);ce=float(-(vt*logits.log_softmax(-1)).sum(-1).mean())
            acc=float((logits.argmax(-1)==vt.argmax(-1)).float().mean())
        row=dict(epoch=epoch,validation_ce=ce,teacher_agreement=acc);history.append(row);print(row,flush=True)
        if ce<best:
            best=ce;best_state=copy.deepcopy(model.state_dict());best_epoch=epoch
model.load_state_dict(best_state);model.eval()
with torch.no_grad():
    p=model(vx,vm).softmax(-1);acc=float((p.argmax(-1)==vt.argmax(-1)).float().mean())
    perm=torch.tensor([2,0,3,1],device=device)
    permuted=model(vx[:,perm],vm[:,perm]).softmax(-1)
    equiv_error=float((permuted-p[:,perm]).abs().max())
assert equiv_error<1e-5
save_file({k:v.cpu().contiguous() for k,v in model.state_dict().items()},out/'structured_policy.safetensors')
config=dict(architecture='Shared option MLP with mean/max set pooling',features=FEATURES,
            parameters=sum(p.numel() for p in model.parameters()),dtype='float32',train_rows=4096,validation_rows=512,
            seed=42,epochs=100,selected_epoch=best_epoch,batch=128,lr=.001,weight_decay=.0001,
            gpu=torch.cuda.get_device_name(),training_seconds=time.monotonic()-start,
            validation_teacher_agreement=acc,validation_ce=best,permutation_max_error=equiv_error,
            note='Standalone structured-feature supervised baseline; not an LLM fine-tune',history=history)
(out/'structured_config.json').write_text(json.dumps(config,indent=2))
print('STRUCTURED_TRAINING_COMPLETE',json.dumps(config),flush=True)