"""Strict BIDS trial reconstruction. Test labels require explicit evaluation=True."""
from pathlib import Path
import json,hashlib
import numpy as np
import pandas as pd
import mne
from scipy.signal import butter,sosfiltfilt,resample_poly
ROOT=Path('datasets/nm000113')
LABELS=['Hello','Help me','Stop','Thank you','Yes']
SPLITS={'train':'00','validation':'01','test':'02'}
def load(subject,split,evaluation=False):
 if split=='test' and not evaluation: raise PermissionError('Test data require explicit final evaluation')
 p=ROOT/f'sub-{subject:02d}'/'eeg'/f'sub-{subject:02d}_task-imaginedSpeech_run-{SPLITS[split]}'
 raw=mne.io.read_raw_edf(str(p)+'_eeg.edf',preload=False,verbose='ERROR')
 ev=pd.read_csv(str(p)+'_events.tsv',sep='\t')
 fs=raw.info['sfreq']; assert fs==256 and len(raw.ch_names)==64
 starts=np.rint(ev.onset.to_numpy()*fs).astype(int); lengths=np.rint(ev.duration.to_numpy()*fs).astype(int)
 assert np.all(lengths==795) and np.array_equal(starts,np.arange(len(ev))*795)
 assert np.array_equal(starts,ev['sample']) and starts[-1]+795<=raw.n_times
 x=np.stack([raw.get_data(start=int(s),stop=int(s+795)) for s in starts])*1e6
 assert np.isfinite(x).all()
 y=np.array([LABELS.index(v) for v in ev.trial_type],dtype=np.int64)
 assert np.array_equal(y+1,ev.value)
 assert len(y)==(300 if split=='train' else 50)
 assert np.array_equal(np.bincount(y,minlength=5),np.full(5,60 if split=='train' else 10))
 return x,y,raw.ch_names

def preprocess(x,indices=None):
 # Official MATLAB t starts at -500 ms; isolate imagined phase [0, 2) seconds.
 # No temporal filtering across trials. Re-reference only measured selected electrodes.
 x=x[:,indices,:] if indices is not None else x
 x=x[:,:,128:640];x=x-x.mean(axis=1,keepdims=True)
 x=sosfiltfilt(butter(4,[1,40],btype='bandpass',fs=256,output='sos'),x,axis=-1)
 return resample_poly(x,25,32,axis=-1).astype(np.float32)

def montage(names,n,kind='distributed'):
 if n==64:return np.arange(64)
 pos=mne.channels.make_standard_montage('standard_1005').get_positions()['ch_pos']
 xyz=np.array([pos[v] for v in names]);xyz/=np.linalg.norm(xyz,axis=1,keepdims=True)
 if kind=='language':
  preferred=['F7','FT7','FC5','T7','C5','CP5','TP7','P7','F5','F3','FC3','C3','CP3','P3','Fz','Cz','F8','FT8','FC6','T8','C6','CP6','TP8','P8','F4','FC4','C4','CP4','P4','Pz','Oz','Fp1']
  ids=[names.index(v) for v in preferred if v in names][:n]
 else:ids=[names.index('Cz')]
 while len(ids)<n:
  d=((xyz[:,None]-xyz[ids][None])**2).sum(-1).min(1);d[ids]=-1;ids.append(int(d.argmax()))
 return np.sort(ids)

def cached(subject,split,evaluation=False):
 if split=='test' and not evaluation:raise PermissionError('test lock')
 p=Path('preprocessing')/f'sub-{subject:02d}_{split}_raw.npz'
 if not p.exists():
  x,y,names=load(subject,split,evaluation);np.savez(p,x=x.astype(np.float32),y=y,names=names)
 z=np.load(p);return z['x'],z['y'],z['names'].tolist()
if __name__=='__main__':
 rows=[];hashes=set()
 for s in range(1,16):
  for split in ['train','validation']:
   x,y,names=cached(s,split)
   for trial in x:
    h=hashlib.sha256(trial.tobytes()).hexdigest();assert h not in hashes,'Duplicate trial';hashes.add(h)
   rows.append(dict(subject=s,split=split,shape=list(x.shape),sfreq=256,unit='microvolt',median_channel_std=float(np.median(x.std(-1))),class_counts=np.bincount(y,minlength=5).tolist()))
 json.dump({'records':rows,'unique_trial_hashes':len(hashes),'test_labels_loaded':False,'montages':{f'{kind}_{n}':[names[i] for i in montage(names,n,kind)] for kind in ['distributed','language'] for n in [64,32,16,8]}},open('reports/bci_validation.json','w'),indent=2)
 print('Validated',len(rows),'records;',len(hashes),'unique train/validation trials',flush=True)