#include <cstdlib>
#include <iostream> 

#include "TChain.h"
#include "TFile.h"
#include "TTree.h"
#include "TString.h"
#include "TObjString.h"
#include "TSystem.h"
#include "TROOT.h"
#include "TPluginManager.h"

#if not defined(__CINT__) || defined(__MAKECINT__)
// needs to be included when makecint runs (ACLIC)
#include "TMVA/Factory.h"
#include "TMVA/Tools.h"
#endif
   
void train_omega2piTMVA(TString myPath = "") 
{

   // this loads the library
   TMVA::Tools::Instance();

   // Create a new root output file
   TString baseName = "omega2pi";
   TString outputFileName = baseName+"_trainTMVA.root";
   TFile* outputFile = TFile::Open(outputFileName, "RECREATE");
   TMVA::Factory *factory = new TMVA::Factory(baseName, outputFile,"!V:!Silent:Color:DrawProgressBar:Transformations=I");

   // 3 pi PID variables 
   //factory->AddVariable("PiPlus1__Timing_FOM := (PiPlus1__ChiSq_Timing_KinFit==PiPlus1__ChiSq_Timing_KinFit) ? TMath::Prob(PiPlus1__ChiSq_Timing_KinFit,PiPlus1__NDF_Timing) : 0.","Timing FOM: #pi+ 1","",'F');
   //factory->AddVariable("PiPlus1__Timing_FOM := (PiPlus1__ChiSq_Timing_Measured==PiPlus1__ChiSq_Timing_Measured) ? TMath::Prob(PiPlus1__ChiSq_Timing_Measured,PiPlus1__NDF_Timing) : 0.","Timing FOM Measured: #pi+ 1","",'F');
   //factory->AddVariable("PiPlus1__DCdEdx_FOM := TMath::Prob(PiPlus1__ChiSq_DCdEdx,PiPlus1__NDF_DCdEdx)","dE/dx FOM: #pi+ 1","",'F');
   factory->AddVariable("PiPlus1__Tracking_FOM := TMath::Prob(PiPlus1__ChiSq_Tracking,PiPlus1__NDF_Tracking)","Tracking FOM: #pi+ 1","",'F');
   //factory->AddVariable("PiPlus1__NDF_Tracking","Track Fit NDF: #pi+ 1","",'I');
   //factory->AddVariable("PiPlus2__Timing_FOM := (PiPlus2__ChiSq_Timing_KinFit==PiPlus2__ChiSq_Timing_KinFit) ? TMath::Prob(PiPlus2__ChiSq_Timing_KinFit,PiPlus2__NDF_Timing) : 0.","Timing FOM: #pi+ 1","",'F');
   //factory->AddVariable("PiPlus2__DCdEdx_FOM := TMath::Prob(PiPlus2__ChiSq_DCdEdx,PiPlus2__NDF_DCdEdx)","dE/dx FOM: #pi+ 2","",'F');
   factory->AddVariable("PiPlus2__Tracking_FOM := TMath::Prob(PiPlus2__ChiSq_Tracking,PiPlus2__NDF_Tracking)","Tracking FOM: #pi+ 2","",'F');
   //factory->AddVariable("PiPlus2__NDF_Tracking","Track Fit NDF: #pi+ 2","",'I');
   //factory->AddVariable("PiMinus1__Timing_FOM := (PiMinus1__ChiSq_Timing_KinFit==PiMinus1__ChiSq_Timing_KinFit) ? TMath::Prob(PiMinus1__ChiSq_Timing_KinFit,PiMinus1__NDF_Timing) : 0.","Timing FOM: #pi+ 1","",'F');
   //factory->AddVariable("PiMinus1__DCdEdx_FOM := TMath::Prob(PiMinus1__ChiSq_DCdEdx,PiMinus1__NDF_DCdEdx)","dE/dx FOM: #pi-","",'F');
   factory->AddVariable("PiMinus1__Tracking_FOM := TMath::Prob(PiMinus1__ChiSq_Tracking,PiMinus1__NDF_Tracking)","Tracking FOM: #pi- 1","",'F');
   //factory->AddVariable("PiMinus1__NDF_Tracking","Track Fit NDF: #pi- 1","",'I');
   //factory->AddVariable("PiMinus2__Timing_FOM := (PiMinus2__ChiSq_Timing_KinFit==PiMinus2__ChiSq_Timing_KinFit) ? TMath::Prob(PiMinus2__ChiSq_Timing_KinFit,PiMinus2__NDF_Timing) : 0.","Timing FOM: #pi+ 2","",'F');
   //factory->AddVariable("PiMinus2__DCdEdx_FOM := TMath::Prob(PiMinus2__ChiSq_DCdEdx,PiMinus2__NDF_DCdEdx)","dE/dx FOM: #pi- 2","",'F');
   factory->AddVariable("PiMinus2__Tracking_FOM := TMath::Prob(PiMinus2__ChiSq_Tracking,PiMinus2__NDF_Tracking)","Tracking FOM: #pi- 2","",'F');
   factory->AddVariable("NumUnusedTracks := NumUnused","Number of Unused Track hypotheses","",'I');
   //factory->AddVariable("PiMinus2__NDF_Tracking","Track Fit NDF: #pi- 2","",'I');
   //factory->AddVariable("Photon1__Timing_FOM := (Photon1__ChiSq_Timing_KinFit==Photon1__ChiSq_Timing_KinFit) ? TMath::Prob(Photon1__ChiSq_Timing_KinFit,Photon1__NDF_Timing) : 0.","Timing FOM: #gamma 1","",'F');
   //factory->AddVariable("Photon2__Timing_FOM := (Photon2__ChiSq_Timing_KinFit==Photon2__ChiSq_Timing_KinFit) ? TMath::Prob(Photon2__ChiSq_Timing_KinFit,Photon2__NDF_Timing) : 0.","Timing FOM: #gamma 2","",'F');
   //factory->AddVariable("Near_Pi0_peak := (TMath::Abs((Photon1__P4_Measured + Photon2__P4_Measured).M() - 135) < 5) ? 1. : 0.","Diphoton mass near #pi^{0}","",'F');
   // Kinematic fit variables
   factory->AddVariable("FOM_KinFit := TMath::Prob(ChiSq_KinFit,NDF_KinFit)","Kinematic Fit Confidence Level","",'F');
   factory->AddVariable("PV_r := MissingProton__X4.Perp()","Missing Proton Primary Vertex Radius","cm",'F');
   factory->AddVariable("MissingProton_PT := MissingProton__P4.Perp()","Missing Proton Transverse Momentum","GeV/c",'F');
   
   // Computed variables from friend tree
   //factory->AddVariable("Unused__Max_Proton_FOM","Unused Proton hypothesis track with largest FOM","",'F');
   factory->AddVariable("Unused__Max_KPlus_FOM","Unused K+ hypothesis track with largest FOM","",'F');
   factory->AddVariable("Unused__Max_KMinus_FOM","Unused K- hypothesis track with largest FOM","",'F');
   //factory->AddVariable("Measured_Missing_Mass := TMath::Abs(Measured__MissingMass)","Measured Missing Mass","",'F');
   //factory->AddVariable("Measured__MissingMass","Measured Missing Mass","",'F');
   //factory->AddVariable("DistanceFromPi0Peak := TMath::Abs(MeasuredPi0Mass_MinusPi0Mass)", "Distance of Diphoton mass from Pi0 peak", "", 'F');
   //factory->AddVariable("MeasuredPi0Mass_MinusPi0Mass", "Distance of Diphoton mass from Pi0 peak", "", 'F');
   //factory->AddVariable("Omega_MeasuredMass", "Measured Omega mass", "", 'F');
   factory->AddVariable("Unused__TotalEnergy_BCAL", "Total unused BCAL energy", "", 'F');
   factory->AddVariable("Unused__TotalEnergy_FCAL", "Total unused FCAL energy", "", 'F');
 
  // load the signal and background event samples from ROOT trees
   TString fnameSignal = "tree_omega2pi.root";
   TString fnameBackground = "tree_omega2pi_background.root";
   TFile *inputSignal = new TFile(myPath+fnameSignal);  
   TFile *inputBackground = new TFile(myPath+fnameBackground);  
   
   TTree *signal     = (TTree*)inputSignal->Get("omega2pi_Tree"); 
   TTree *background_combinatoric = (TTree*)inputSignal->Get("omega2pi_Tree");
   TTree *background = (TTree*)inputBackground->Get("omega2piBackground_Tree");
   
   // friend tree with trueSignal and other computed variables
   signal->AddFriend("omega2piFriend_Tree","tree_omega2piFriend.root");
   background_combinatoric->AddFriend("omega2piFriend_Tree","tree_omega2piFriend.root");
   background->AddFriend("omega2piFriend_Tree","tree_omega2piFriendBackground.root");
   
   // global event weights per tree (see below for setting event-wise weights)
   Double_t signalWeight     = 1.0; 
   Double_t backgroundWeight = 1.0;
   
   factory->AddSignalTree    ( signal,     signalWeight     );
   factory->AddBackgroundTree( background_combinatoric, signalWeight );
   factory->AddBackgroundTree( background, backgroundWeight );

   TString bothCut = ""; //" && PiPlus1__NDF_Tracking>5 && PiPlus2__NDF_Tracking>5 && PiMinus1__NDF_Tracking>5 && PiMinus2__NDF_Tracking>5";
   TString signalCut = "TrueSignal"; 
   TString backgroundCut = "!TrueSignal"; 
   signalCut+=bothCut; backgroundCut+=bothCut;

   // Apply additional cuts on the signal and background samples (can be different)
   TCut mycuts = signalCut;
   TCut mycutb = backgroundCut;

   // tell the factory to use all remaining events in the trees after training for testing:
   factory->PrepareTrainingAndTestTree( mycuts, mycutb,"nTrain_Signal=0:nTrain_Background=250000:nTest_Signal=0:nTest_Background=2000000:SplitMode=Random:!V" );

	   // book MVA method
   factory->BookMethod( TMVA::Types::kBDT, "BDT","!H:!V:MaxDepth=2:VarTransform=G,D,G,D" );
   //factory->BookMethod( TMVA::Types::kCuts, "OptimizedCuts","");
   // --------------------------------------------------------------------------------------------------

   // ---- Now you can tell the factory to train, test, and evaluate the MVAs

   // Train MVAs using the set of training events
   factory->TrainAllMethods();

   // ---- Evaluate all MVAs using the set of test events
   factory->TestAllMethods();

   // ----- Evaluate and compare performance of all configured MVAs
   factory->EvaluateAllMethods();    

   // --------------------------------------------------------------
   
   // Save the output
   outputFile->Close();
   delete factory;

}