#include "TMVA/QuickMVAProbEstimator.h"
 
#include <iostream>

#include "TMath.h"

#ifndef ROOT_TMVA_MsgLogger
#include "TMVA/MsgLogger.h"
#endif


void TMVA::QuickMVAProbEstimator::AddEvent(Double_t val, Double_t weight, Int_t type){
   EventInfo ev;
   ev.eventValue=val; ev.eventWeight=weight; ev.eventType=type; 
  
   fEvtVector.push_back(ev);
   if (fIsSorted) fIsSorted=false;
                       
}

 
Double_t TMVA::QuickMVAProbEstimator::GetMVAProbAt(Double_t value){
   // Well.. if it's fast is actually another question all together, merely
   // it's a quick and dirty simple kNN approach to the 1-Dim signal/backgr. MVA
   // distributions. 


   if (!fIsSorted) {
      std::sort(fEvtVector.begin(),fEvtVector.end(),TMVA::QuickMVAProbEstimator::compare), fIsSorted=true;
   }

   Double_t     percentage = 0.1;
   UInt_t  nRange = TMath::Max(fNMin,(UInt_t) (fEvtVector.size() * percentage));
   nRange = TMath::Min(fNMax,nRange);
   // just make sure that nRange > you total number of events
   if (nRange > fEvtVector.size()) {
      nRange = fEvtVector.size()/3.;
      Log() << kWARNING  << " !!  you have only " << fEvtVector.size() << " of events.. . I choose "
            << nRange << " for the quick and dirty kNN MVAProb estimate" << Endl;
   }

   EventInfo tmp; tmp.eventValue=value;
   std::vector<EventInfo>::iterator it = std::upper_bound(fEvtVector.begin(),fEvtVector.end(),tmp,TMVA::QuickMVAProbEstimator::compare);
  
   UInt_t iLeft=0, iRight=0;
   Double_t nSignal=0;
   Double_t nBackgr=0;

   while ( (iLeft+iRight) < nRange){
      if ( fEvtVector.end() > it+iRight+1){
         iRight++;
         if ( ((it+iRight))->eventType == 0) nSignal+=((it+iRight))->eventWeight;
         else                                nBackgr+=((it+iRight))->eventWeight;
      }
      if ( fEvtVector.begin() <= it-iLeft-1){
         iLeft++;
         if ( ((it-iLeft))->eventType == 0) nSignal+=((it-iLeft))->eventWeight;
         else                               nBackgr+=((it-iLeft))->eventWeight;
      }    
   }

   Double_t mvaProb = (nSignal+nBackgr) ? nSignal/(nSignal+nBackgr) : -1 ;
   return mvaProb;

}
