"""Block-held-out semantic classification, excluding EOG/stim and reading data."""
import os
os.environ.setdefault('MPLCONFIGDIR',os.path.abspath('.cache/matplotlib'))
import argparse,json,re,time,warnings,joblib
from pathlib import Path
import mne,numpy as np
from scipy.signal import welch
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 threadpoolctl import threadpool_limits
threadpool_limits(1)
ROOT=Path('datasets/chisco/derivatives/preprocessed_fif/sub-01/eeg');MAP=json.load(open('datasets/chisco_metadata/json/textmaps.json'))
def block(p):return int(re.search(r'run-(\d+)',p.name)[1])
def extract(split,evaluation=False):
 if split=='test' and not evaluation:raise PermissionError('Chisco final test locked')
 cache=Path('preprocessing')/f'chisco_sub01_{split}.npz'
 if cache.exists():return dict(np.load(cache))
 allowed={'train':range(1,28),'validation':range(28,37),'test':range(37,46)}[split]
 xx=[];yy=[];gg=[];words=[];eogs=[];records=[]
 for p in sorted(ROOT.glob('*task-imagine*.fif'),key=block):
  if block(p) not in allowed:continue
  with warnings.catch_warnings():
   warnings.simplefilter('ignore');e=mne.read_epochs(p,preload=True,verbose='ERROR')
  assert e.info['sfreq']==500 and e.tmin==5 and e.tmax==8.3
  labels=np.array([MAP.get(w.strip(),-1) for w in e.metadata.Word]);valid=labels>=0
  eeg=e.copy().pick('eeg');assert len(eeg.ch_names)==122
  # Omit 300 ms after visual cue offset; use imagined-only late window.
  x=eeg.get_data(copy=True)[valid,:,-1500:]*1e6;x-=x.mean(axis=1,keepdims=True)
  f,power=welch(x,fs=500,nperseg=500,axis=-1)
  feat=np.concatenate([np.log(power[:,:,(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)
  eo=e.copy().pick('eog').get_data(copy=True)[valid,:,-1500:]*1e6
  _,ep=welch(eo,fs=500,nperseg=500,axis=-1);ef=np.concatenate([np.log(ep[:,:,(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)
  xx.append(feat);eogs.append(ef);yy.extend(labels[valid]);gg.extend([block(p)]*valid.sum());words.extend(e.metadata.Word.to_numpy()[valid]);records.append({'block':block(p),'trials':int(valid.sum()),'discarded_unmapped':int((~valid).sum())})
 result={'x':np.concatenate(xx),'eog':np.concatenate(eogs),'y':np.array(yy),'block':np.array(gg),'text':np.array(words)}
 np.savez(cache,**result);json.dump(records,open(f'reports/chisco_{split}_blocks.json','w'),indent=2);return result

def metrics(y,p):return dict(accuracy=accuracy_score(y,p),balanced_accuracy=balanced_accuracy_score(y,p),macro_f1=f1_score(y,p,labels=range(39),average='macro',zero_division=0),confusion_matrix=confusion_matrix(y,p,labels=range(39)).tolist())
def train():
 a=extract('train');b=extract('validation');assert a['y'].min()==0 and a['y'].max()==38
 # Remove shared sentences from training using validation text identity only.
 keep=~np.isin(a['text'],b['text']);results=[]
 for kind in ['x','eog']:
  best=None;search=[];t=time.time()
  for c in [.001,.01,.1,1]:
   m=make_pipeline(StandardScaler(),LogisticRegression(C=c,max_iter=1500,class_weight='balanced'));m.fit(a[kind][keep],a['y'][keep]);met=metrics(b['y'],m.predict(b[kind]));search.append({'C':c,**met})
   if best is None or met['balanced_accuracy']>best[0]:best=(met['balanced_accuracy'],m,c,met)
  name=f'chisco_sub01_{kind}_blocksplit';joblib.dump(best[1],f'checkpoints/{name}.joblib')
  r={'name':name,'split':'validation','metrics':best[3],'C':best[2],'search':search,'classes':39,'uniform_chance':1/39,'train_n':int(keep.sum()),'validation_n':len(b['y']),'train_classes':np.bincount(a['y'][keep],minlength=39).tolist(),'validation_classes':np.bincount(b['y'],minlength=39).tolist(),'removed_validation_text_overlap':int((~keep).sum()),'seconds':time.time()-t,'interpretation':'EEG semantic classifier' if kind=='x' else 'EOG artifact control'}
  json.dump(r,open(f'experiments/{name}.json','w'),indent=2);print(name,best[3]['accuracy'],best[0],flush=True)
if __name__=='__main__':
 from experiment_guard import assert_tuning_open
 assert_tuning_open()
 train()