"""Multiband supervised CSP, fit only on training trials."""
import os
os.environ.setdefault('MPLCONFIGDIR',os.path.abspath('.cache/matplotlib'))
from pathlib import Path
import time,json,joblib,numpy as np,mne
from scipy.signal import butter,sosfiltfilt
from sklearn.base import BaseEstimator,TransformerMixin
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from bci_data import cached,preprocess,montage
from classical import metrics
mne.set_log_level('ERROR')
class FilterBankCSP(BaseEstimator,TransformerMixin):
 def __init__(self,n_components=6):self.n_components=n_components
 def filter(self,x,a,b):return sosfiltfilt(butter(4,[a,b],btype='bandpass',fs=200,output='sos'),x,axis=-1)
 def fit(self,x,y):
  self.filters_=[]
  for a,b in [(4,8),(8,13),(13,20),(20,30),(30,40)]:
   c=mne.decoding.CSP(n_components=min(self.n_components,x.shape[1]-1),reg='oas',log=True,cov_est='concat',norm_trace=False,rank='full');c.fit(self.filter(x,a,b),y);self.filters_.append((a,b,c))
  return self
 def transform(self,x):return np.concatenate([c.transform(self.filter(x,a,b)) for a,b,c in self.filters_],axis=1)

def run():
 for n in [64,32,16,8]:
  for s in range(1,16):
   name=f'strict_fbcsp_distributed_{n}_sub-{s:02d}';out=Path('experiments')/(name+'.json')
   if out.exists():continue
   x,y,names=cached(s,'train');v,z,_=cached(s,'validation');idx=montage(names,n);x=preprocess(x,idx).astype(float);v=preprocess(v,idx).astype(float);t=time.time()
   csp=FilterBankCSP().fit(x,y);a=csp.transform(x);b=csp.transform(v);best=None
   for c in [.001,.01,.1,1,10]:
    m=make_pipeline(StandardScaler(),LogisticRegression(C=c,max_iter=1500,class_weight='balanced'));m.fit(a,y);met=metrics(z,m.predict(b))
    if best is None or met['balanced_accuracy']>best[0]:best=(met['balanced_accuracy'],m,c,met)
   joblib.dump({'model':make_pipeline(csp,best[1]),'kind':'fbcsp','indices':idx,'names':names,'subject':s},f'checkpoints/{name}.joblib');out.write_text(json.dumps({'name':name,'model':'fbcsp','channels':n,'montage':'distributed','subject':s,'split':'validation','metrics':best[3],'C':best[2],'seconds':time.time()-t},indent=2));print(name,best[0],flush=True)
if __name__=='__main__':
 from experiment_guard import assert_tuning_open
 assert_tuning_open()
 # Stable pickle module name for checkpoints saved by CLI execution.
 from fbcsp import run
 run()