"""Calibrate/evaluate frozen candidates. No training or test threshold selection."""
import argparse,json,hashlib
from pathlib import Path
import numpy as np
from irstd.core import candidate_pool,image_stats,aggregate,calibrate,components
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('--gate');p.add_argument('--split',choices=['dev','cal','test'],default='test')
 p.add_argument('--calibrate',action='store_true');p.add_argument('--certificate');p.add_argument('--threshold',type=float,default=.5)
 p.add_argument('--base-threshold',type=float,default=.1);p.add_argument('--cap',type=int,default=64);p.add_argument('--radius',type=float,default=3.)
 p.add_argument('--alpha',type=float,default=.1);p.add_argument('--delta',type=float,default=.05)
 p.add_argument('--risk',choices=['count','any'],default='count');p.add_argument('--pixel-baseline',action='store_true')
 p.add_argument('--calibration-method',choices=['bonferroni','fixed_sequence'],default='bonferroni')
 a=p.parse_args()
 if a.calibrate and (a.split!='cal' or a.pixel_baseline or a.certificate):p.error('Calibration requires cal split, fixed candidates and no certificate input')
 if a.pixel_baseline and (a.gate or a.certificate):p.error('Pixel baseline cannot use gate/candidate certificate')
 if not 0<=a.threshold<=1:p.error('threshold outside [0,1]')
 rows=read_manifest(a.manifest);root=Path(a.predictions)
 provenance_path=root/'provenance.json'
 if not provenance_path.exists():raise ValueError('Missing prediction provenance.json')
 provenance=json.loads(provenance_path.read_text())
 if provenance.get('manifest_sha256')!=sha(a.manifest):raise ValueError('Prediction manifest differs')
 pipeline=dict(manifest_sha256=sha(a.manifest),provenance=provenance,
               base_threshold=a.base_threshold,cap=a.cap,radius=a.radius,
               gate_sha256=sha(a.gate) if a.gate else None,
               prediction_hashes={r['id']:sha(root/(r['id']+'.npy')) for r in rows},
               input_hashes={r['id']:dict(image=sha(r['image']),mask=sha(r['mask'])) for r in rows},
               source_hashes={name:sha(Path(__file__).resolve().parents[1]/name) for name in ['irstd/core.py','irstd/models.py','scripts/evaluate.py']},
               evaluation_version=1)
 fingerprint=hashlib.sha256(json.dumps(pipeline,sort_keys=True).encode()).hexdigest()
 gate=None
 if a.gate:
  import torch
  from irstd.models import CandidateGate
  ck=torch.load(a.gate,map_location='cpu',weights_only=True);cfg=ck['config']
  for k in ('cap','base_threshold','radius'):
   if cfg[k]!=getattr(a,k):raise ValueError('Gate pipeline differs: '+k)
  if cfg['manifest_sha256']!=sha(a.manifest):raise ValueError('Gate manifest differs')
  gate=CandidateGate();gate.load_state_dict(ck['state_dict']);gate.eval()
 thresholds=np.linspace(0,1,101).tolist() if a.calibrate else [a.threshold]
 reject_all=False
 if a.certificate:
  cert=json.loads(Path(a.certificate).read_text())
  if cert['pipeline_fingerprint']!=fingerprint:raise ValueError('Certificate invalidated by changed pipeline/data')
  reject_all=cert['reject_all'];thresholds=[cert['threshold'] if not reject_all else 1.]
 selected=[r for r in rows if r['split']==a.split]
 if not selected:raise ValueError('Empty split')
 fp=[];stats=[];per_image=[]
 for r in selected:
  prob=np.load(root/(r['id']+'.npy'),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)
  scores=features[:,0]
  if gate is not None and len(features):
   with torch.no_grad():scores=gate(torch.from_numpy(features)).sigmoid().numpy()
  row_stats=[]
  for t in thresholds:
   keep=[] if reject_all else [m for m,s in zip(masks,scores) if s>=t]
   if a.pixel_baseline:keep=components(prob>=t)
   row_stats.append(image_stats(keep,gt,a.radius))
  fp.append([s['fp'] for s in row_stats]);stats.append(row_stats[0])
  per_image.append(dict(id=r['id'],group=r['group'],**row_stats[0]))
 if a.calibrate:
  result=calibrate(fp,thresholds,a.cap,a.alpha,a.delta,a.risk,a.calibration_method)
  result.update(pipeline_fingerprint=fingerprint,pipeline=pipeline,
                calibration_ids=[r['id'] for r in selected],
                calibration_groups=len(set(r['group'] for r in selected)),
                assumption_notice='Conditional on iid images; group labels alone do not verify independence. Do not pool adjacent video frames.')
 else:
  result=dict(metrics=aggregate(stats),per_image=per_image,pipeline_fingerprint=fingerprint,
              split=a.split,threshold=None if reject_all else thresholds[0],reject_all=reject_all,
              certificate_used=a.certificate,mode='pixel' if a.pixel_baseline else 'fixed_candidate')
 dest=Path(a.out);dest.parent.mkdir(parents=True,exist_ok=True)
 if dest.exists():raise FileExistsError(dest)
 dest.write_text(json.dumps(result,indent=2,allow_nan=False));print(dest)
if __name__=='__main__':main()