import os
os.environ.setdefault('OPENBLAS_NUM_THREADS','1');os.environ.setdefault('OMP_NUM_THREADS','1')
import argparse,json,time,joblib
from pathlib import Path
import numpy as np
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score,balanced_accuracy_score,f1_score,confusion_matrix
from pyriemann.estimation import Covariances
from pyriemann.tangentspace import TangentSpace
from pyriemann.classification import MDM
from scipy.signal import welch
from bci_data import cached,preprocess,montage
from threadpoolctl import threadpool_limits
threadpool_limits(1)
def metrics(y,p):
 return dict(accuracy=accuracy_score(y,p),balanced_accuracy=balanced_accuracy_score(y,p),macro_f1=f1_score(y,p,average='macro',zero_division=0),confusion_matrix=confusion_matrix(y,p,labels=range(5)).tolist())
def features(x,kind):
 if kind in ['ts','mdm']:return Covariances(estimator='oas').fit_transform(x)
 f,p=welch(x,fs=200,nperseg=200,axis=-1)
 return np.concatenate([np.log(p[:,:,(f>=a)&(f<b)].mean(-1)+1e-12) for a,b in [(1,4),(4,8),(8,13),(13,20),(20,30),(30,40)]],axis=1)
def estimator(kind,c):
 if kind=='mdm':return MDM(metric='logeuclid')
 parts=[TangentSpace(metric='riemann')] if kind=='ts' else []
 return make_pipeline(*parts,StandardScaler(),LogisticRegression(C=c,max_iter=1500,class_weight='balanced'))
def run(subject,n,kind,selection):
 name=f'strict_{kind}_{selection}_{n}_sub-{subject:02d}';out=Path('experiments')/(name+'.json')
 if out.exists():return
 t=time.time();x,y,names=cached(subject,'train');v,vy,_=cached(subject,'validation');idx=montage(names,n,selection)
 a=features(preprocess(x,idx),kind);b=features(preprocess(v,idx),kind)
 best=None;scores=[]
 for c in ([1] if kind=='mdm' else [.001,.01,.1,1,10]):
  model=estimator(kind,c);model.fit(a,y);pred=model.predict(b);m=metrics(vy,pred);scores.append({'C':c,**m})
  if best is None or m['balanced_accuracy']>best[0]:best=(m['balanced_accuracy'],model,c,m)
 joblib.dump({'model':best[1],'indices':idx,'kind':kind,'names':names,'subject':subject,'preprocessing':'car_1-40_imagined0to2s_200hz'},Path('checkpoints')/(name+'.joblib'))
 out.write_text(json.dumps(dict(name=name,subject=subject,channels=n,montage=selection,model=kind,split='validation',selected_C=best[2],metrics=best[3],search=scores,seconds=time.time()-t),indent=2))
 print(name,round(best[0],4),round(time.time()-t,1),flush=True)
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)));p.add_argument('--models',nargs='+',default=['bandpower','mdm','ts']);p.add_argument('--montage',default='distributed');args=p.parse_args()
 for n in args.channels:
  for kind in args.models:
   for s in args.subjects:run(s,n,kind,args.montage)