import argparse,json,random,os
from pathlib import Path
os.environ.setdefault('CUBLAS_WORKSPACE_CONFIG', ':4096:8')
import numpy as np
import torch
from torch.nn import functional as F
from irstd.models import CandidateGate,gate_loss
from irstd.core import candidate_pool,canonical_labels
from irstd.data import read_manifest,read_pair,sha

def main():
 p=argparse.ArgumentParser();p.add_argument('--manifest',required=True);p.add_argument('--predictions',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('--mode',choices=['bce','constrained'],default='constrained');p.add_argument('--no-size',action='store_true')
 p.add_argument('--alpha-train',type=float,default=.1);p.add_argument('--cap',type=int,default=64);p.add_argument('--base-threshold',type=float,default=.1)
 p.add_argument('--radius',type=float,default=3.);a=p.parse_args()
 if a.epochs<1 or a.alpha_train<0: p.error('Invalid epochs/budget')
 random.seed(a.seed);np.random.seed(a.seed);torch.manual_seed(a.seed);torch.use_deterministic_algorithms(True)
 rows=read_manifest(a.manifest);cached=[];files={}
 for r in rows:
  if r['split']!='train':continue
  path=Path(a.predictions)/(r['id']+'.npy');prob=np.load(path,allow_pickle=False)
  _,gt=read_pair(r)
  if prob.shape!=gt.shape:raise ValueError('Prediction shape mismatch')
  masks,features=candidate_pool(prob,a.base_threshold,a.cap);labels,areas=canonical_labels(masks,gt,a.radius)
  cached.append(tuple(torch.from_numpy(v) for v in (features,labels,areas)));files[r['id']]=sha(path)
 if not cached:raise ValueError('No train images')
 gate=CandidateGate();opt=torch.optim.Adam(gate.parameters(),lr=1e-3);dual=0.
 out=Path(a.out);out.mkdir(parents=True,exist_ok=False)
 config=vars(a)|dict(manifest_sha256=sha(a.manifest),training_prediction_hashes=files)
 # Fixed epochs: tune this on dev in separate runs, never on calibration/test labels.
 with (out/'gate.jsonl').open('w') as f:
  for epoch in range(a.epochs):
   random.shuffle(cached);upper_total=0.;loss_total=0.
   for x,y,area in cached:
    if not len(x):continue  # zero surrogate contribution; retain image in denominator
    opt.zero_grad();s=gate(x)
    if a.mode=='constrained':
     loss,upper=gate_loss(s,y,area,dual,a.alpha_train,not a.no_size)
    else:
     loss=F.binary_cross_entropy_with_logits(s,y);_,upper=gate_loss(s,y,area,0.,a.alpha_train)
    loss.backward();opt.step();upper_total+=float(upper.detach());loss_total+=float(loss.detach())
   mean_upper=upper_total/len(cached)
   if a.mode=='constrained':dual=max(0.,dual+.01*(mean_upper-a.alpha_train))
   f.write(json.dumps(dict(epoch=epoch,loss=loss_total/len(cached),surrogate=mean_upper,dual=dual))+'\n');f.flush()
 torch.save(dict(state_dict=gate.state_dict(),config=config),out/'gate.pt')
 (out/'config.json').write_text(json.dumps(config,indent=2))
if __name__=='__main__':main()