"""Inference accepts measured EEG only; labels never enter the prediction method."""
import os,time
os.environ.setdefault('MPLCONFIGDIR',os.path.abspath('.cache/matplotlib'))
from pathlib import Path
import numpy as np,joblib
from bci_data import preprocess,LABELS
from classical import features
class Predictor:
 def __init__(self,path):
  self.path=str(path);self.neural=Path(path).suffix=='.pt';self.backbone=None
  if self.neural:
   import torch
   from neural import build
   torch.set_num_threads(2);self.state=torch.load(path,map_location='cpu',weights_only=False);self.model=build(self.state['kind'],self.state['channels'],[self.state['names'][i] for i in self.state['indices']]);self.model.load_state_dict(self.state['model']);self.model.eval()
  else:
   self.state=joblib.load(path);self.model=self.state['model']
 def predict(self,raw):
  t=time.perf_counter();idx=self.state['indices'];x=preprocess(np.asarray(raw)[None],idx)
  if self.neural:
   import torch
   x=np.clip(x/self.state['scale'],-20,20)
   if self.state['kind'].startswith('cbramod_'):x=x.reshape(-1,len(idx),2,200)
   with torch.inference_mode():prob=self.model(torch.tensor(x)).softmax(-1).numpy()[0]
   label=int(prob.argmax())
  else:
   kind=self.state['kind']
   if kind in ['reve','csbrain']:
    import torch
    from foundation_models import Foundation
    if self.backbone is None:self.backbone=Foundation(kind,[self.state['names'][i] for i in idx]).eval();torch.set_num_threads(2)
    with torch.inference_mode():f=self.backbone(torch.tensor(x/self.state['scale'])).float().numpy()
   elif kind=='cbramod':
    import torch
    from cbramod_probe import backbone
    if self.backbone is None:self.backbone=backbone().eval();torch.set_num_threads(2)
    with torch.inference_mode():f=self.backbone(torch.tensor(x/100).reshape(-1,len(idx),2,200)).flatten(1).numpy()
   elif self.state.get('mode')=='late_matched':
    from late_control import late
    f=features(late(np.asarray(raw)[None].copy()),'bandpower')
   elif self.state.get('mode')=='precue':
    from controls import precue
    f=features(precue(np.asarray(raw)[None]),'bandpower')
   else:f=x.astype(float) if kind in ['fbcsp','fbts'] else features(x,kind)
   prob=self.model.predict_proba(f)[0];label=int(self.model.predict(f)[0])
  return {'prediction':label,'command':LABELS[label],'confidence':float(prob[label]),'probabilities':dict(zip(LABELS,map(float,prob))),'latency_ms':(time.perf_counter()-t)*1000,'measured_channels':len(idx),'model':Path(self.path).name}