import os,sys,time,json,hashlib,argparse
from pathlib import Path
os.environ.setdefault('MPLCONFIGDIR',os.path.abspath('.cache/matplotlib'))
sys.path.insert(0,os.path.abspath('external/CBraMod'))
import torch,numpy as np,joblib
from torch import nn
from bci_data import cached,preprocess,montage
from classical import metrics
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from models.cbramod import CBraMod

def backbone():
 m=CBraMod();state=torch.load('models/cbramod_pretrained.pth',map_location='cpu',weights_only=True);m.load_state_dict(state,strict=True);m.proj_out=nn.Identity();return m

def get_features(m,subject,split,n,device,evaluation=False):
 p=Path('preprocessing')/f'strict_cbramod_{n}_sub-{subject:02d}_{split}.npy'
 x,y,names=cached(subject,split,evaluation)
 idx=montage(names,n)
 if p.exists():return np.load(p),y,idx
 x=preprocess(x,idx)/100 # µV / 100 per official speech dataset loader
 feats=[];m.eval()
 with torch.inference_mode():
  for a in torch.tensor(x).split(16):feats.append(m(a.to(device).reshape(-1,n,2,200)).flatten(1).cpu().numpy())
 f=np.concatenate(feats);np.save(p,f);return f,y,idx
if __name__=='__main__':
 from experiment_guard import assert_tuning_open
 assert_tuning_open()
 p=argparse.ArgumentParser();p.add_argument('--channels',nargs='+',type=int,default=[64]);p.add_argument('--subjects',nargs='+',type=int,default=list(range(1,16)));a=p.parse_args()
 torch.set_num_threads(4);device='cuda' if torch.cuda.is_available() else 'cpu';m=backbone().to(device)
 print('strict official checkpoint loaded',sum(p.numel() for p in m.parameters()),device,flush=True)
 for n in a.channels:
  for s in a.subjects:
   name=f'strict_cbramod_distributed_{n}_sub-{s:02d}';out=Path('experiments')/(name+'.json')
   if out.exists():continue
   from neural import expired
   if expired():sys.exit(0)
   t=time.time();x,y,idx=get_features(m,s,'train',n,device);v,z,_=get_features(m,s,'validation',n,device);best=None;search=[]
   for c in [.0001,.001,.01,.1,1]:
    clf=make_pipeline(StandardScaler(),LogisticRegression(C=c,max_iter=1000,class_weight='balanced'));clf.fit(x,y);met=metrics(z,clf.predict(v));search.append({'C':c,**met})
    if best is None or met['balanced_accuracy']>best[0]:best=(met['balanced_accuracy'],clf,c,met)
   joblib.dump({'model':best[1],'kind':'cbramod','indices':idx,'subject':s,'preprocessing':'strict_0to2s_uVdiv100','checkpoint_sha256':hashlib.sha256(Path('models/cbramod_pretrained.pth').read_bytes()).hexdigest()},Path('checkpoints')/(name+'.joblib'))
   out.write_text(json.dumps({'name':name,'model':'cbramod','channels':n,'montage':'distributed','subject':s,'split':'validation','metrics':best[3],'C':best[2],'search':search,'seconds':time.time()-t,'device':device,'mode':'frozen_flattened_features_linear_probe'},indent=2));print(name,best[0],time.time()-t,flush=True)