"""Filter-bank Riemannian tangent-space classifier; every reference mean is train-only."""
import os
os.environ.setdefault('MPLCONFIGDIR',os.path.abspath('.cache/matplotlib'))
from pathlib import Path
import json,time,joblib,numpy as np
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 pyriemann.estimation import Covariances
from pyriemann.tangentspace import TangentSpace
from bci_data import cached,preprocess,montage
from classical import metrics
class FilterBankTangent(BaseEstimator,TransformerMixin):
 def cov(self,x,a,b):
  z=sosfiltfilt(butter(4,[a,b],btype='bandpass',fs=200,output='sos'),x,axis=-1)
  return Covariances(estimator='oas').transform(z)
 def fit(self,x,y=None):
  self.spaces_=[]
  for a,b in [(1,4),(4,8),(8,13),(13,20),(20,30),(30,40)]:self.spaces_.append((a,b,TangentSpace(metric='riemann').fit(self.cov(x,a,b))))
  return self
 def transform(self,x):return np.concatenate([t.transform(self.cov(x,a,b)) for a,b,t in self.spaces_],axis=1)
def run():
 from neural import expired
 for n in [32,16,8,64]:
  for s in range(1,16):
   if expired():return
   name=f'strict_fbts_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();fb=FilterBankTangent().fit(x);a=fb.transform(x);b=fb.transform(v);best=None;search=[]
   for c in [.0001,.001,.01,.1,1]:
    m=make_pipeline(StandardScaler(),LogisticRegression(C=c,max_iter=1500,class_weight='balanced'));m.fit(a,y);met=metrics(z,m.predict(b));search.append({'C':c,**met})
    if best is None or met['balanced_accuracy']>best[0]:best=(met['balanced_accuracy'],m,c,met)
   joblib.dump({'model':make_pipeline(fb,best[1]),'kind':'fbts','indices':idx,'names':names,'subject':s},f'checkpoints/{name}.joblib');out.write_text(json.dumps({'name':name,'model':'fbts','channels':n,'montage':'distributed','subject':s,'split':'validation','metrics':best[3],'C':best[2],'search':search,'seconds':time.time()-t},indent=2));print(name,best[0],time.time()-t,flush=True)
if __name__=='__main__':
 from experiment_guard import assert_tuning_open
 assert_tuning_open()
 from fbts import run
 run()