% train a classification model 
% using current observations for Features 1-3
% exclude SWPC events from training set

clear all

pwd_dir=pwd;
method_name = 'RobustBoost';
ens_num=20;

station_list={'ott','iqa','frn'};

root = ['~/geospace_work/'];

for station = station_list
 
  cd([root char(station)])
  disp(['Training ' char(station)])

 filename = [char(station) '_train_test_sets.mat']
 load(filename)
 
 Xtrain=table2array(T_train(:,2:end-1));
 ytrain=table2array(T_train(:,end));

 [i j]=find(isinf(Xtrain));
 Xtrain(i,:)=[]; ytrain(i)=[];
 [i j]=find(isnan(Xtrain));
 Xtrain(i,:)=[];ytrain(i)=[];

f=find(ytrain>5e4); Xtrain(f,:)=[];ytrain(f)=[];
f=find(log10(ytrain)<-2);Xtrain(f,:)=[];ytrain(f)=[];

 load([char(station) '_thresholds.mat'])

% C=pca(Xtrain);
%Xtrain=Xtrain*C;

for class = 1:5
  disp(['Training class ' num2str(class)])
  clear X y t

  P=randperm(size(Xtrain,1));
  Xtrain = Xtrain(P,:);
  ytrain = ytrain(P);
  
  thre = threshold_list(class);

  pos=find(ytrain/60>=thre);
  num_pos = length(pos);
  neg =find(ytrain/60<thre);
  num_neg = length(neg);

  if num_neg > num_pos
    neg(num_pos+1:end)=[];
  else
    pos(num_neg+1:end)=[];
  end
  
  q = quantile(ytrain(neg),0.2);
  f=find(ytrain(neg)<q);neg(f)=[];

% create Train set
  X = [Xtrain(pos,:); Xtrain(neg,:)];
  y = [ytrain(pos); ytrain(neg)];
  t = double(y/60>=thre);
    
   tic
   tmp=fitcensemble(X,t,'Nprint',2,'NumLearningCycles',ens_num,'Method',method_name,'ScoreTransform','logit');
   Mdl{class} = compact(tmp);
   toc
   disp(' ')
     
  end
  
filename = [char(station) '_trained_model_classification_grey_box.mat']
save(filename,'Mdl')

end

cd(pwd_dir)
