"""CUDA-only foundation adaptation. Official train/validation only; resumable epochs."""
import os,json,time,argparse,random
from pathlib import Path
from recovery_state import completed, atomic_json
import numpy as np,torch,joblib
from torch import nn
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from bci_data import cached,preprocess,montage
from classical import metrics
from foundation_models import Foundation,FoundationClassifier
from gpu_utils import require_cuda,expired,atomic_save,cpu_state,install_signals

def scale_for(kind):return 1000. if kind=='reve' else 100.
def tensors(subject,n,kind):
 x,y,names=cached(subject,'train');v,z,_=cached(subject,'validation');idx=montage(names,n);scale=scale_for(kind)
 return preprocess(x,idx)/scale,y,preprocess(v,idx)/scale,z,idx,names

def probe(kind,n,s):
 name=f'strict_{kind}_probe_distributed_{n}_sub-{s:02d}';out=Path('experiments')/(name+'.json')
 if completed(name):return
 device=require_cuda();x,y,v,z,idx,names=tensors(s,n,kind);encoder=Foundation(kind,[names[i] for i in idx]).to(device).eval();features=[];start=time.time()
 for split,a in [('train',x),('validation',v)]:
  path=Path('preprocessing')/f'{name}_{split}_features.npy'
  if path.exists():features.append(np.load(path));continue
  chunks=[]
  with torch.inference_mode():
   for q in torch.tensor(a).split(16):
    if expired():return
    with torch.autocast('cuda',dtype=torch.bfloat16):f=encoder(q.to(device))
    chunks.append(f.float().cpu().numpy())
  f=np.concatenate(chunks);np.save(path,f);features.append(f)
 del encoder;torch.cuda.empty_cache();best=None;search=[]
 for c in [.001,.01,.1,1.]:
  if expired():return
  model=make_pipeline(StandardScaler(),LogisticRegression(C=c,max_iter=1500,class_weight='balanced'));model.fit(features[0],y);m=metrics(z,model.predict(features[1]));search.append({'C':c,**m})
  if best is None or m['balanced_accuracy']>best[0]:best=(m['balanced_accuracy'],model,c,m)
 Path('reports/validation_predictions').mkdir(exist_ok=True)
 np.savez_compressed(f'reports/validation_predictions/{name}.npz',y=z,pred=best[1].predict(features[1]))
 joblib.dump({'model':best[1],'kind':kind,'indices':idx,'names':names,'subject':s,'scale':scale_for(kind),'preprocessing':'car_1-40_imagined0to2s_200hz','feature_precision':'cuda_bfloat16'},f'checkpoints/{name}.joblib')
 atomic_json(out, {'name':name,'model':kind+'_probe','subject':s,'channels':n,'montage':'distributed','split':'validation','selected_C':best[2],'metrics':best[3],'search':search,'seconds':time.time()-start,'device':'cuda','precision':'bfloat16'});print(name,best[3],flush=True)

def finetune(kind,mode,n,s,epochs=50,lr=.0001,seed=2026):
 name=f'strict_{kind}_{mode}_distributed_{n}_sub-{s:02d}_seed-{seed}';out=Path('experiments')/(name+'.json')
 if completed(name):return
 device=require_cuda();random.seed(seed);np.random.seed(seed);torch.manual_seed(seed);torch.cuda.manual_seed_all(seed);starttime=time.time()
 x,y,v,z,idx,names=tensors(s,n,kind);x=np.clip(x,-20,20);v=np.clip(v,-20,20);model=FoundationClassifier(kind,[names[i] for i in idx],mode).to(device)
 opt=torch.optim.AdamW([{'params':[p for p in model.encoder.parameters() if p.requires_grad],'lr':lr/10},{'params':model.head.parameters(),'lr':lr}],weight_decay=.01)
 loader=torch.utils.data.DataLoader(torch.utils.data.TensorDataset(torch.tensor(x),torch.tensor(y)),batch_size=16,shuffle=True);vx=torch.tensor(v,device=device)
 best=-1;best_epoch=-1;start=0;history=[];last=Path(f'checkpoints/{name}_last.pt');bestpath=Path(f'checkpoints/{name}.pt')
 if last.exists():
  state=torch.load(last,map_location='cpu',weights_only=False);model.load_state_dict(state['model']);opt.load_state_dict(state['optimizer']);start=state['epoch']+1;best=state['best'];best_epoch=state['best_epoch'];history=state['history'];torch.set_rng_state(state['rng']);torch.cuda.set_rng_state_all(state['cuda_rng']);random.setstate(state['python_rng']);np.random.set_state(state['numpy_rng'])
 criterion=nn.CrossEntropyLoss(label_smoothing=.05)

 if not last.exists() and bestpath.exists():
  state=torch.load(bestpath,map_location='cpu',weights_only=False);model.load_state_dict(state['model']);best=state['validation_metrics']['balanced_accuracy'];best_epoch=state['epoch'];start=best_epoch+1
  print(name,'Resuming best weights only: optimizer/RNG snapshot unavailable',flush=True)
 if last.exists() and start-1-best_epoch>=12:start=epochs
 for epoch in range(start,epochs):
  if expired():return
  model.train();losses=[]
  for a,b in loader:
   a=a.to(device);b=b.to(device);a=a*(.9+.2*torch.rand((len(a),1,1),device=device))+torch.randn_like(a)*(.02 if kind=='csbrain' else .002)
   opt.zero_grad(set_to_none=True)
   with torch.autocast('cuda',dtype=torch.bfloat16):logits=model(a);loss=criterion(logits,b)
   if not torch.isfinite(loss):raise FloatingPointError('Nonfinite training loss')
   loss.backward();nn.utils.clip_grad_norm_(model.parameters(),1.);opt.step();losses.append(loss.item())
  model.eval()
  with torch.inference_mode(),torch.autocast('cuda',dtype=torch.bfloat16):pred=torch.cat([model(q) for q in vx.split(16)]).float().argmax(1).cpu().numpy()
  m=metrics(z,pred);history.append({'epoch':epoch,'loss':float(np.mean(losses)),**m});score=m['balanced_accuracy']
  if score>best:
   best=score;best_epoch=epoch;atomic_save({'model':cpu_state(model),'kind':kind+'_'+mode,'channels':n,'indices':idx,'names':names,'scale':np.full((1,n,1),scale_for(kind),dtype=np.float32),'seed':seed,'subjects':[s],'epoch':epoch,'validation_metrics':m,'training_device':'cuda','precision':'bfloat16'},bestpath)
  atomic_save({'model':cpu_state(model),'optimizer':opt.state_dict(),'epoch':epoch,'best':best,'best_epoch':best_epoch,'history':history,'rng':torch.get_rng_state(),'cuda_rng':torch.cuda.get_rng_state_all(),'python_rng':random.getstate(),'numpy_rng':np.random.get_state()},last)
  print(name,epoch,score,float(np.mean(losses)),flush=True)
  if epoch-best_epoch>=12:break
 if expired():return
 beststate=torch.load(bestpath,map_location='cpu',weights_only=False)
 model.load_state_dict(beststate['model']);model.eval()
 with torch.inference_mode(),torch.autocast('cuda',dtype=torch.bfloat16):pred=torch.cat([model(q) for q in vx.split(16)]).float().argmax(1).cpu().numpy()
 Path('reports/validation_predictions').mkdir(exist_ok=True)
 np.savez_compressed(f'reports/validation_predictions/{name}.npz',y=z,pred=pred)
 atomic_json(out, {'name':name,'model':kind+'_'+mode,'channels':n,'montage':'distributed','subjects':[s],'seed':seed,'split':'validation','best_epoch':best_epoch,'metrics':beststate['validation_metrics'],'history':history,'seconds':time.time()-starttime,'device':'cuda','precision':'bfloat16','lr':lr,'parameters':sum(p.numel() for p in model.parameters())})

if __name__=='__main__':
 from experiment_guard import assert_tuning_open
 assert_tuning_open();require_cuda();install_signals();torch.set_num_threads(4)
 p=argparse.ArgumentParser();p.add_argument('--model',choices=['reve','csbrain'],required=True);p.add_argument('--mode',choices=['probe','partial','full'],default='probe');p.add_argument('--channels',nargs='+',type=int,default=[64,32,16,8]);p.add_argument('--subjects',nargs='+',type=int,default=list(range(1,16)));p.add_argument('--epochs',type=int,default=50);a=p.parse_args()
 for n in a.channels:
  for s in a.subjects:
   if expired():raise SystemExit(3)
   probe(a.model,n,s) if a.mode=='probe' else finetune(a.model,a.mode,n,s,a.epochs)