import os
os.environ.setdefault('MPLCONFIGDIR',os.path.abspath('.cache/matplotlib'))
import json,time,joblib
from pathlib import Path
import numpy as np
from scipy.signal import resample_poly,butter,sosfiltfilt
from sklearn.feature_selection import f_classif
from classical import features,estimator,metrics
from bci_data import cached,preprocess,montage

def ranks(x,y):
 # Univariate training-only channel ranking; no validation signal/labels used.
 f=features(preprocess(x),'bandpower').reshape(len(y),6,64)
 score=np.array([np.nan_to_num(f_classif(f[:,:,i],y)[0],nan=0).mean() for i in range(64)])
 return np.argsort(score)[::-1]
def run_ranked():
 for s in range(1,16):
  x,y,names=cached(s,'train');v,z,_=cached(s,'validation');rank=ranks(x,y)
  for n in [32,16,8]:
   name=f'strict_bandpower_trainrank_{n}_sub-{s:02d}';out=Path('experiments')/(name+'.json')
   if out.exists():continue
   idx=np.sort(rank[:n]);a=features(preprocess(x,idx),'bandpower');b=features(preprocess(v,idx),'bandpower');best=None
   for c in [.001,.01,.1,1,10]:
    m=estimator('bandpower',c);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':best[1],'indices':idx,'kind':'bandpower','names':names,'subject':s},f'checkpoints/{name}.joblib');out.write_text(json.dumps({'name':name,'model':'bandpower','montage':'trainrank','channels':n,'subject':s,'split':'validation','metrics':best[3],'selected_C':best[2],'ranked_channels':[names[i] for i in rank]},indent=2))
   print(name,best[0],flush=True)
def precue(x):
 x=x[:,:,:128];x=x-x.mean(1,keepdims=True)
 return resample_poly(sosfiltfilt(butter(4,[4,40],btype='bandpass',fs=256,output='sos'),x,axis=-1),25,32,axis=-1).astype(np.float32)
def run_controls():
 for s in range(1,16):
  x,y,names=cached(s,'train');v,z,_=cached(s,'validation')
  for mode in ['precue','permuted']:
   name=f'control_{mode}_sub-{s:02d}';out=Path('experiments')/(name+'.json')
   if out.exists():continue
   a=features(precue(x) if mode=='precue' else preprocess(x),'bandpower');b=features(precue(v) if mode=='precue' else preprocess(v),'bandpower')
   yy=np.random.default_rng(2026+s).permutation(y) if mode=='permuted' else y
   m=estimator('bandpower',.01);m.fit(a,yy);met=metrics(z,m.predict(b))
   joblib.dump({'model':m,'kind':'bandpower','mode':mode,'indices':np.arange(64),'subject':s},f'checkpoints/{name}.joblib');out.write_text(json.dumps({'name':name,'model':mode,'channels':64,'subject':s,'split':'validation','metrics':met,'C':.01},indent=2));print(name,met['accuracy'],flush=True)
def run_loso():
 # Each outer subject completely absent from fit and hyperparameter selection.
 data={}
 for s in range(1,16):
  data[s]={}
  for split in ['train','validation']:
   x,y,_=cached(s,split);data[s][split]=(features(preprocess(x),'bandpower'),y)
 for held in range(1,16):
  name=f'loso_bandpower_64_sub-{held:02d}';out=Path('experiments')/(name+'.json')
  if out.exists():continue
  others=[s for s in data if s!=held];x=np.concatenate([data[s]['train'][0] for s in others]);y=np.concatenate([data[s]['train'][1] for s in others]);v=np.concatenate([data[s]['validation'][0] for s in others]);z=np.concatenate([data[s]['validation'][1] for s in others]);best=None
  for c in [.001,.01,.1,1]:
   m=estimator('bandpower',c);m.fit(x,y);met=metrics(z,m.predict(v))
   if best is None or met['balanced_accuracy']>best[0]:best=(met['balanced_accuracy'],m,c,met)
  joblib.dump({'model':best[1],'kind':'bandpower','indices':np.arange(64),'subject':held,'excluded_subject':held},f'checkpoints/{name}.joblib');out.write_text(json.dumps({'name':name,'model':'loso_bandpower','channels':64,'subject':held,'split':'inner_validation_other_subjects','metrics':best[3],'C':best[2],'training_subjects':others},indent=2));print(name,best[0],flush=True)
if __name__=='__main__':
 from experiment_guard import assert_tuning_open
 assert_tuning_open()
 run_ranked();run_controls();run_loso()