"""Train an EEGNet on your own preprocessed epochs, then import experiment.json.
Usage: python train_custom.py epochs.npz --test-group session_3
NPZ: X float [N,3,384] in microvolts; y integers 0/1; groups strings [N].
X must be 4–38 Hz filtered epochs, C3/Cz/C4 in that order, 128 Hz, 3 s.
Prepare continuous recordings separately per session; never filter across a split.
Only de-identified data you have permission to process. No data is uploaded.
"""
import argparse,base64,hashlib,json,tempfile
from pathlib import Path
import numpy as np
def require(condition, message):
    # Input and export guarantees must also hold when Python is run with -O.
    if not condition:
        raise ValueError(message)

def prepare_epochs(X, y, groups, test_group):
    X=np.asarray(X,dtype='float32');y=np.asarray(y);groups=np.asarray(groups)
    require(X.ndim==3 and X.shape[1:]==(3,384), 'X must be [trials,3,384]')
    require(y.shape==groups.shape==(len(X),), 'One label and recording-group ID per trial')
    require(np.isfinite(X).all() and set(np.unique(y))=={0,1}, 'Finite samples and binary labels required')
    require(groups.dtype.kind in 'US', 'Recording-group IDs must be strings, with no missing values')
    groups=groups.astype(str)
    require(all(g.strip() for g in groups), 'Every trial needs a nonempty recording-group ID')
    require(isinstance(test_group,str) and bool(test_group.strip()), 'Choose a nonempty test-group ID')
    test=groups==test_group
    require(2<=test.sum()<=300 and (~test).sum()>=20, 'Need 2-300 test trials and at least 20 train trials')
    require(len(np.unique(y[test]))==len(np.unique(y[~test]))==2, 'Both splits need both classes')
    # Every trial from the named recording group is held out before fitting anything.
    trainX=X[~test];mu=trainX.mean(axis=(0,2),keepdims=True);sd=trainX.std(axis=(0,2),keepdims=True).clip(1e-6)
    require(np.isfinite(mu).all() and np.isfinite(sd).all(), 'Training values exceed supported numerical range; check microvolt units')
    trainX=(trainX-mu)/sd;testX=np.round((X[test]-mu)/sd,6).astype('float32')
    require(np.isfinite(trainX).all() and np.isfinite(testX).all(), 'Normalization produced nonfinite samples; check units and preprocessing')
    require(abs(testX).max()<=100, 'Extreme normalized values; inspect artifact contamination')
    return trainX,y[~test].astype('int64'),testX,y[test].astype('int64'),mu,sd

def main(argv=None):
    p=argparse.ArgumentParser(description=__doc__)
    p.add_argument('npz');p.add_argument('--test-group',required=True)
    p.add_argument('--output',default='experiment.json')
    p.add_argument('--class-0',default='Class 0');p.add_argument('--class-1',default='Class 1')
    a=p.parse_args(argv)
    require(all(name.strip() and len(name)<=60 for name in [a.class_0,a.class_1]), 'Class names must contain 1-60 characters')
    require(a.class_0!=a.class_1, 'Use distinct names for the two classes')
    with np.load(a.npz,allow_pickle=False) as d:
        require(all(key in d for key in ['X','y','groups']), 'NPZ must contain X, y and groups')
        trainX,trainY,testX,testY,mu,sd=prepare_epochs(d['X'],d['y'],d['groups'],a.test_group)
    import torch
    from braindecode.models import EEGNet
    import onnxruntime as ort
    torch.set_num_threads(2);torch.manual_seed(42)
    model=EEGNet(n_chans=3,n_outputs=2,n_times=384,sfreq=128)
    opt=torch.optim.Adam(model.parameters(),lr=.001)
    tx,ty=torch.from_numpy(trainX),torch.from_numpy(trainY)
    for epoch in range(60):
        model.train()
        for idx in torch.randperm(len(ty)).split(8):
            opt.zero_grad();loss=torch.nn.functional.cross_entropy(model(tx[idx]),ty[idx]);loss.backward();opt.step()
    model.eval()
    with tempfile.TemporaryDirectory() as tmp:
        f=Path(tmp)/'model.onnx'
        torch.onnx.export(model,torch.from_numpy(testX[:1]),str(f),input_names=['eeg'],output_names=['logits'],dynamic_axes={'eeg':{0:'batch'},'logits':{0:'batch'}},opset_version=17,dynamo=False)
        with torch.no_grad(): expected=model(torch.from_numpy(testX)).numpy()
        observed=ort.InferenceSession(str(f),providers=['CPUExecutionProvider']).run(None,{'eeg':testX})[0]
        require(np.isfinite(expected).all() and np.isfinite(observed).all() and np.max(abs(expected-observed))<1e-4,'ONNX parity failed')
        weights=f.read_bytes()
    output={'format':'daybreak-experiment-v1','model_name':'EEGNet','model_base64':base64.b64encode(weights).decode(),'model_sha256':hashlib.sha256(weights).hexdigest(),'classes':[a.class_0,a.class_1],'protocol':f'User-provided preprocessed epochs. Hold out group {a.test_group[:100]}; all other groups train. Train-only channel mean/std. EEGNet, Adam lr=.001,60 epochs,seed42. No tuning on test set.','normalization_mean_uv':mu.flatten().tolist(),'normalization_std_uv':sd.flatten().tolist(),'data':{'schema_version':1,'channels':['C3','Cz','C4'],'sfreq':128,'trials':testX.tolist(),'labels':testY.tolist()}}
    text=json.dumps(output,separators=(',',':'),allow_nan=False)
    require(len(text.encode('utf-8'))<=12*1024*1024,'Too large: export fewer test trials')
    Path(a.output).write_text(text,encoding='utf-8')
    print(f'Saved {a.output}. Test accuracy: {(observed.argmax(1)==testY).mean():.3f}. This is not independent validation if you used these test outcomes to select a model.')

if __name__=='__main__':
    main()