"""Single-device, batch-one reference training; no dataset auto-download."""
import argparse,json,random,os
from pathlib import Path
os.environ.setdefault('CUBLAS_WORKSPACE_CONFIG', ':4096:8')
import numpy as np
import torch
from irstd.models import TinyUNet,segmentation_loss
from irstd.data import read_manifest,read_pair,sha

def main():
 p=argparse.ArgumentParser(); p.add_argument('--manifest',required=True);p.add_argument('--out',required=True)
 p.add_argument('--epochs',type=int,default=100);p.add_argument('--seed',type=int,default=0)
 p.add_argument('--loss',choices=['bce','dice','bce_dice','focal','iou'],default='bce_dice')
 p.add_argument('--device',default='cuda'); a=p.parse_args()
 if a.epochs<1: p.error('epochs must be positive')
 random.seed(a.seed);np.random.seed(a.seed);torch.manual_seed(a.seed)
 torch.use_deterministic_algorithms(True);torch.backends.cudnn.benchmark=False
 rows=read_manifest(a.manifest);tr=[r for r in rows if r['split']=='train'];dv=[r for r in rows if r['split']=='dev']
 if not tr or not dv: raise ValueError('Train and development splits required')
 out=Path(a.out);out.mkdir(parents=True,exist_ok=False)
 model=TinyUNet().to(a.device);opt=torch.optim.AdamW(model.parameters(),lr=1e-3,weight_decay=1e-4)
 config=vars(a)|dict(manifest_sha256=sha(a.manifest),torch=str(torch.__version__),numpy=np.__version__,cuda=torch.version.cuda)
 (out/'config.json').write_text(json.dumps(config,indent=2));best=float('inf')
 def batch(r):
  x,y=read_pair(r)
  return torch.from_numpy(x)[None,None].to(a.device),torch.from_numpy(y.astype(np.float32))[None,None].to(a.device)
 with (out/'train.jsonl').open('w') as f:
  for epoch in range(a.epochs):
   model.train();random.shuffle(tr);total=0.
   for r in tr:
    x,y=batch(r);opt.zero_grad();loss=segmentation_loss(model(x),y,a.loss)
    loss.backward();opt.step();total+=loss.item()
   model.eval();dev=0.
   with torch.no_grad():
    for r in dv:
     x,y=batch(r);dev+=segmentation_loss(model(x),y,a.loss).item()
   dev/=len(dv);f.write(json.dumps(dict(epoch=epoch,train_loss=total/len(tr),dev_loss=dev))+'\n');f.flush()
   if dev<best:
    best=dev;torch.save(dict(state_dict=model.state_dict(),config=config,epoch=epoch),out/'best.pt')
if __name__=='__main__':main()