"""Freeze choices BEFORE opening test labels; no post-test tuning permitted."""
import os
os.environ.setdefault('MPLCONFIGDIR',os.path.abspath('.cache/matplotlib'))
import json,hashlib,datetime,time,collections,argparse
from pathlib import Path
import numpy as np,joblib,torch
from bci_data import cached,preprocess,LABELS
from classical import features,metrics
from neural import build

def candidates():
 rows=[]
 for p in Path('experiments').glob('strict_*.json'):
  r=json.load(open(p))
  ext='.pt' if 'seed' in r else '.joblib';cp=Path('checkpoints')/(r['name']+ext)
  if not cp.exists():continue
  if len(r.get('subjects',[]))>1:
   for s in r['subjects']:
    rr=dict(r);rr['name']=r['name']+f'_eval-sub-{s:02d}';rr['model']=r['model']+'_pooled';rr['subject']=s;rr['metrics']=r['subject_validation_metrics'][str(s)];rows.append((s,rr,cp))
  else:
   s=r.get('subject',r.get('subjects',[None])[0]);rows.append((s,r,cp))
 return rows

def freeze():
 if Path('experiments/coordination/FINAL_TEST_SEALED_BY_USER').exists():raise RuntimeError('User-sealed final test: review validation audit before authorizing final evaluation')
 if Path('experiments/coordination/HOLD_FINAL_EVALUATION').exists():raise RuntimeError('GPU experiments pending: final evaluation is on hold')
 path=Path('reports/frozen_selection.json')
 if path.exists():raise RuntimeError('Selection already frozen; never overwrite after test evaluation')
 rows=candidates();groups=collections.defaultdict(list)
 for s,r,p in rows:groups[(r['model'],r['montage'],r['channels'],r.get('seed','deterministic'))].append((s,r,p))
 complete={k:v for k,v in groups.items() if len(v)==15}
 bestkey=max(sorted(complete,key=str),key=lambda k:np.mean([r['metrics']['balanced_accuracy'] for _,r,_ in complete[k]]))
 demo={str(s):str(p) for s,r,p in complete[bestkey]}
 personal={str(s):str(max(sorted([(r,p) for ss,r,p in rows if ss==s],key=lambda x:str(x[1])),key=lambda x:(x[0]['metrics']['balanced_accuracy'],x[0]['metrics']['macro_f1'],-x[0]['channels']))[1]) for s in range(1,16)}
 controls=[]
 for pattern in ['control_*.json','loso_*.json']:
  for p in Path('experiments').glob(pattern):
   r=json.load(open(p));cp=Path('checkpoints')/(r['name']+'.joblib')
   if cp.exists():controls.append((r['subject'],r,cp))
 manifest={'frozen_utc':datetime.datetime.now(datetime.timezone.utc).isoformat(),'primary_family':list(bestkey),'primary_mean_validation_balanced_accuracy':float(np.mean([r['metrics']['balanced_accuracy'] for _,r,_ in complete[bestkey]])),'demo_subject_models':demo,'personalized_selection':personal,'evaluations':[{'subject':s,'record':r,'checkpoint':str(p),'sha256':hashlib.sha256(p.read_bytes()).hexdigest()} for s,r,p in rows+controls],'chisco_checkpoints':{str(p):hashlib.sha256(p.read_bytes()).hexdigest() for p in Path('checkpoints').glob('chisco_*') if '_last' not in p.name},'rule':'Primary family maximizes mean validation balanced accuracy across all 15 subjects. Per-subject hyperparameters fixed by validation. Secondary subject-specific architecture selection recorded before test, with ties broken by macro-F1, fewer channels and lexicographic checkpoint order. No retuning after freeze.'}
 chisco_candidates=[]
 for kind in ['x','eegnet','cbramod']:
  p=Path(f'experiments/chisco_sub01_{kind}_blocksplit.json')
  if p.exists():
   record=json.load(open(p));chisco_candidates.append((record['metrics']['balanced_accuracy'],kind))
 if chisco_candidates:
  score,kind=max(chisco_candidates);manifest['chisco_primary_model']={'model':kind,'validation_balanced_accuracy':score,'rule':'Highest validation balanced accuracy among EEG-only pipelines; excludes EOG and prior controls'}
 subset_candidates=[]
 for kind in ['x','cbramod']:
  p=Path(f'experiments/chisco_sub01_five_{kind}_blocksplit.json')
  if p.exists():
   record=json.load(open(p));subset_candidates.append((record['metrics']['balanced_accuracy'],kind,record['class_ids']))
 if subset_candidates:
  score,kind,class_ids=max(subset_candidates);manifest['chisco_five_primary_model']={'model':kind,'class_ids':class_ids,'validation_balanced_accuracy':score,'rule':'Five classes chosen by training frequency; EEG pipeline selected by validation balanced accuracy'}
 reference=[]
 for p in Path('experiments').glob('reference3s_*.json'):
  record=json.load(open(p));cp=Path('checkpoints')/(record['name']+'.pt')
  reference.append({'record':record,'checkpoint':str(cp),'sha256':hashlib.sha256(cp.read_bytes()).hexdigest()})
 manifest['reference_evaluations']=reference
 if reference:manifest['reference_primary']=max(reference,key=lambda r:r['record']['metrics']['balanced_accuracy'])['record']['name']
 manifest['source_hashes']={str(p):hashlib.sha256(p.read_bytes()).hexdigest() for p in list(Path('scripts').glob('*.py'))+list(Path('configs').glob('*.json'))+[Path('requirements.lock.txt'),Path('datasets/nemar_manifest.json')]}
 path.write_text(json.dumps(manifest,indent=2));print('FROZEN',bestkey,manifest['primary_mean_validation_balanced_accuracy'],flush=True)

def predictions(cp,x):
 if cp.suffix=='.pt':
  state=torch.load(cp,map_location='cpu',weights_only=False);m=build(state['kind'],state['channels']);m.load_state_dict(state['model']);m.eval();a=preprocess(x,state['indices']);a=np.clip(a/state['scale'],-20,20)
  if state['kind'].startswith('cbramod_'):a=a.reshape(-1,state['channels'],2,200)
  with torch.inference_mode():prob=torch.cat([m(q).softmax(-1) for q in torch.tensor(a).split(16)]).numpy()
  return prob.argmax(1),prob
 state=joblib.load(cp);a=preprocess(x,state['indices']);kind=state['kind']
 if kind in ['reve','csbrain']:
  from foundation_models import Foundation
  m=Foundation(kind,[state['names'][i] for i in state['indices']]).eval()
  with torch.inference_mode():f=torch.cat([m(q) for q in torch.tensor(a/state['scale']).split(16)]).float().numpy()
 elif kind=='cbramod':
  from cbramod_probe import backbone
  m=backbone().eval()
  with torch.inference_mode():f=torch.cat([m(q).flatten(1) for q in torch.tensor(a/100).reshape(-1,len(state['indices']),2,200).split(16)]).numpy()
 elif state.get('mode')=='late_matched':
  from late_control import late
  f=features(late(x.copy()),'bandpower')
 elif state.get('mode')=='precue':
  from controls import precue
  f=features(precue(x),'bandpower')
 else:f=a.astype(float) if kind in ['fbcsp','fbts'] else features(a,kind)
 return state['model'].predict(f),state['model'].predict_proba(f)

def evaluate():
 if Path('experiments/coordination/FINAL_TEST_SEALED_BY_USER').exists():raise RuntimeError('User-sealed final test: no evaluation permitted')
 torch.set_num_threads(2);manifest=json.load(open('reports/frozen_selection.json'));Path('reports/predictions').mkdir(exist_ok=True);out=[]
 from verify_official_test import verify
 verify()
 auditpath=Path('reports/test_evaluation_audit.json')
 if not auditpath.exists():auditpath.write_text(json.dumps({'started_utc':datetime.datetime.now(datetime.timezone.utc).isoformat(),'selection_frozen_utc':manifest['frozen_utc'],'selection_sha256':hashlib.sha256(Path('reports/frozen_selection.json').read_bytes()).hexdigest(),'mode':'Final test; all tuning scripts refuse to run after freeze'},indent=2))
 for entry in manifest['evaluations']:
  r=entry['record'];cp=Path(entry['checkpoint']);s=entry['subject'];p=Path('reports/predictions')/(r['name']+'.npz')
  assert hashlib.sha256(cp.read_bytes()).hexdigest()==entry['sha256'],'Checkpoint changed after freeze'
  if p.exists():z=np.load(p);y=z['y'];pred=z['pred'];prob=z['prob'];latency=float(z['batch_ms'])
  else:
   x,y,_=cached(s,'test',evaluation=True);t=time.perf_counter();pred,prob=predictions(cp,x);latency=1000*(time.perf_counter()-t);np.savez(p,y=y,pred=pred,prob=prob,batch_ms=latency)
  rr={'name':r['name'],'subject':s,'model':r['model'],'channels':r['channels'],'montage':r.get('montage','control'),'seed':r.get('seed','deterministic'),'metrics':metrics(y,pred),'checkpoint_bytes':cp.stat().st_size,'batch_ms_including_load':latency,'checkpoint':str(cp)};out.append(rr);print(r['name'],rr['metrics']['accuracy'],flush=True)
 Path('reports/test_subject_results.json').write_text(json.dumps(out,indent=2))
 audit=json.load(open(auditpath));audit.update(completed_utc=datetime.datetime.now(datetime.timezone.utc).isoformat(),participant_configuration_results=len(out));auditpath.write_text(json.dumps(audit,indent=2))
 # Assemble validation-selected primary and secondary systems from frozen paths.
 bypath={(r['checkpoint'],r['subject']):r for r in out}
 for key in ['demo_subject_models','personalized_selection']:
  selected=[bypath[(manifest[key][str(s)],s)] for s in range(1,16)];Path(f'reports/{key}_test.json').write_text(json.dumps(selected,indent=2))

if __name__=='__main__':
 p=argparse.ArgumentParser();p.add_argument('action',choices=['freeze','test']);a=p.parse_args();freeze() if a.action=='freeze' else evaluate()