#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");

   // 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__FOM_Tracking := TMath::Prob(PiPlus1__ChiSq_Tracking,PiPlus1__NDF_Tracking)","Track Fit FOM: #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__FOM_Tracking := TMath::Prob(PiPlus2__ChiSq_Tracking,PiPlus2__NDF_Tracking)","Track Fit FOM: #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__FOM_Tracking := TMath::Prob(PiMinus1__ChiSq_Tracking,PiMinus1__NDF_Tracking)","Track Fit FOM: #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__FOM_Tracking := TMath::Prob(PiMinus2__ChiSq_Tracking,PiMinus2__NDF_Tracking)","Track Fit FOM: #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("Proton__Timing_FOM := (Proton__ChiSq_Timing_KinFit==Proton__ChiSq_Timing_KinFit) ? TMath::Prob(Proton__ChiSq_Timing_KinFit,Proton__NDF_Timing) : 0.","Timing FOM: Proton","",'F');
   factory->AddVariable("Proton__Timing_FOM := (Proton__ChiSq_Timing_Measured==Proton__ChiSq_Timing_Measured) ? TMath::Prob(Proton__ChiSq_Timing_Measured,Proton__NDF_Timing) : 0.","Timing FOM Measured: Proton","",'F');
   factory->AddVariable("Proton__DCdEdx_FOM := TMath::Prob(Proton__ChiSq_DCdEdx,Proton__NDF_DCdEdx)","dE/dx FOM: Proton 1","",'F');
   //factory->AddVariable("Proton__FOM_Tracking := TMath::Prob(Proton__ChiSq_Tracking,Proton__NDF_Tracking)","Track Fit FOM: Proton","",'I');
   //factory->AddVariable("PV_r := Proton__X4_Measured.Perp()","Measured Primary Vertex Radius","cm",'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("NumUnusedTracks := NumUnused", "Number of unused track hypotheses", "", 'I');
   //factory->AddVariable("PV_r := MissingNeutron__X4.Perp()","Primary Vertex Radius","cm",'F');
   //factory->AddVariable("MissingNeutron_PT := MissingNeutron__P4.Perp()","Missing Neutron Transverse Momentum","GeV/c",'F');
   
   // These variables were suggested by Paul for the training
   //factory->AddVariable("Proton_dEdx_CDC := (Proton__dEdx_CDC == Proton__dEdx_CDC) ? Proton__dEdx_CDC : 0.","Proton dEdX CDC","",'F');
   //factory->AddVariable("Proton_dEdx_FDC := (Proton__dEdx_FDC == Proton__dEdx_FDC) ? Proton__dEdx_FDC : 0.","Proton dEdX FDC","",'F');
   //factory->AddVariable("Proton_dEdx_TOF := (Proton__dEdx_TOF == Proton__dEdx_TOF) ? Proton__dEdx_TOF : 0.","Proton dEdX TOF","",'F');
   //factory->AddVariable("Proton_dEdx_ST := (Proton__dEdx_ST == Proton__dEdx_ST) ? Proton__dEdx_ST : 0.","Proton dEdX ST","",'F');
   factory->AddVariable("Proton_dEdx_Total := ((Proton__dEdx_CDC + Proton__dEdx_FDC + Proton__dEdx_TOF + Proton__dEdx_ST) == (Proton__dEdx_CDC + Proton__dEdx_FDC + Proton__dEdx_TOF + Proton__dEdx_ST)) ? (Proton__dEdx_CDC + Proton__dEdx_FDC + Proton__dEdx_TOF + Proton__dEdx_ST) : 0.","Proton dEdX Total","",'F');
   //factory->AddVariable("Proton_TrackDOCAToCalorimeterShower := (Proton__TrackDOCA_BCAL < Proton__TrackDOCA_FCAL) ? Proton__TrackDOCA_BCAL : Proton__TrackDOCA_FCAL", "Closest calorimeter shower to Proton track", "", 'F');
   //factory->AddVariable("Proton_TotalCalorimeterEnergy := Proton__Energy_BCAL + Proton__Energy_FCAL","Proton total calorimeter energy","",'F');

   //factory->AddVariable("PiMinus1_dEdx_CDC := (PiMinus1__dEdx_CDC == PiMinus1__dEdx_CDC) ? PiMinus1__dEdx_CDC : 0.","#pi-1 dEdX CDC","",'F');
   //factory->AddVariable("PiMinus1_dEdx_FDC := (PiMinus1__dEdx_FDC == PiMinus1__dEdx_FDC) ? PiMinus1__dEdx_FDC : 0.","#pi-1 dEdX FDC","",'F');
   //factory->AddVariable("PiMinus1_dEdx_TOF := (PiMinus1__dEdx_TOF == PiMinus1__dEdx_TOF) ? PiMinus1__dEdx_TOF : 0.","#pi-1 dEdX TOF","",'F');
   //factory->AddVariable("PiMinus1_dEdx_ST := (PiMinus1__dEdx_ST == PiMinus1__dEdx_ST) ? PiMinus1__dEdx_ST : 0.","#pi-1 dEdX ST","",'F');
   factory->AddVariable("PiMinus1_dEdx_Total := ((PiMinus1__dEdx_CDC + PiMinus1__dEdx_FDC + PiMinus1__dEdx_TOF + PiMinus1__dEdx_ST) == (PiMinus1__dEdx_CDC + PiMinus1__dEdx_FDC + PiMinus1__dEdx_TOF + PiMinus1__dEdx_ST)) ? (PiMinus1__dEdx_CDC + PiMinus1__dEdx_FDC + PiMinus1__dEdx_TOF + PiMinus1__dEdx_ST) : 0.","#pi-1 dEdX Total","",'F');
   //factory->AddVariable("PiMinus1_TrackDOCAToCalorimeterShower := (PiMinus1__TrackDOCA_BCAL < PiMinus1__TrackDOCA_FCAL) ? PiMinus1__TrackDOCA_BCAL : PiMinus1__TrackDOCA_FCAL", "Closest calorimeter shower to #pi-1 track", "", 'F');
   //factory->AddVariable("PiMinus1_TotalCalorimeterEnergy := PiMinus1__Energy_BCAL + PiMinus1__Energy_FCAL","#pi-1 total calorimeter energy","",'F');

   //factory->AddVariable("PiMinus2_dEdx_CDC := (PiMinus2__dEdx_CDC == PiMinus2__dEdx_CDC) ? PiMinus2__dEdx_CDC : 0.","#pi-2 dEdX CDC","",'F');
   //factory->AddVariable("PiMinus2_dEdx_FDC := (PiMinus2__dEdx_FDC == PiMinus2__dEdx_FDC) ? PiMinus2__dEdx_FDC : 0.","#pi-2 dEdX FDC","",'F');
   //factory->AddVariable("PiMinus2_dEdx_TOF := (PiMinus2__dEdx_TOF == PiMinus2__dEdx_TOF) ? PiMinus2__dEdx_TOF : 0.","#pi-2 dEdX TOF","",'F');
   //factory->AddVariable("PiMinus2_dEdx_ST := (PiMinus2__dEdx_ST == PiMinus2__dEdx_ST) ? PiMinus2__dEdx_ST : 0.","#pi-2 dEdX ST","",'F');
   //factory->AddVariable("PiMinus2_TrackDOCAToCalorimeterShower := (PiMinus2__TrackDOCA_BCAL < PiMinus2__TrackDOCA_FCAL) ? PiMinus2__TrackDOCA_BCAL : PiMinus2__TrackDOCA_FCAL", "Closest calorimeter shower to #pi-2 track", "", 'F');
   //factory->AddVariable("PiMinus2_TotalCalorimeterEnergy := PiMinus2__Energy_BCAL + PiMinus2__Energy_FCAL","#pi-2 total calorimeter energy","",'F');

   //factory->AddVariable("PiPlus1_dEdx_CDC := (PiPlus1__dEdx_CDC == PiPlus1__dEdx_CDC) ? PiPlus1__dEdx_CDC : 0.","#pi+1 dEdX CDC","",'F');
   //factory->AddVariable("PiPlus1_dEdx_FDC := (PiPlus1__dEdx_FDC == PiPlus1__dEdx_FDC) ? PiPlus1__dEdx_FDC : 0.","#pi+1 dEdX FDC","",'F');
   //factory->AddVariable("PiPlus1_dEdx_TOF := (PiPlus1__dEdx_TOF == PiPlus1__dEdx_TOF) ? PiPlus1__dEdx_TOF : 0.","#pi+1 dEdX TOF","",'F');
   //factory->AddVariable("PiPlus1_dEdx_ST := (PiPlus1__dEdx_ST == PiPlus1__dEdx_ST) ? PiPlus1__dEdx_ST : 0.","#pi+1 dEdX ST","",'F');
   factory->AddVariable("PiPlus1_dEdx_Total := ((PiPlus1__dEdx_CDC + PiPlus1__dEdx_FDC + PiPlus1__dEdx_TOF + PiPlus1__dEdx_ST) == (PiPlus1__dEdx_CDC + PiPlus1__dEdx_FDC + PiPlus1__dEdx_TOF + PiPlus1__dEdx_ST)) ? (PiPlus1__dEdx_CDC + PiPlus1__dEdx_FDC + PiPlus1__dEdx_TOF + PiPlus1__dEdx_ST) : 0.","#pi+1 dEdX Total","",'F');
   //factory->AddVariable("PiPlus1_TrackDOCAToCalorimeterShower := (PiPlus1__TrackDOCA_BCAL < PiPlus1__TrackDOCA_FCAL) ? PiPlus1__TrackDOCA_BCAL : PiPlus1__TrackDOCA_FCAL", "Closest calorimeter shower to #pi+1 track", "", 'F');
   //factory->AddVariable("PiPlus1_TotalCalorimeterEnergy := PiPlus1__Energy_BCAL + PiPlus1__Energy_FCAL","#pi+1 total calorimeter energy","",'F');

   //factory->AddVariable("PiPlus2_dEdx_CDC := (PiPlus2__dEdx_CDC == PiPlus2__dEdx_CDC) ? PiPlus2__dEdx_CDC : 0.","#pi+2 dEdX CDC","",'F');
   //factory->AddVariable("PiPlus2_dEdx_FDC := (PiPlus2__dEdx_FDC == PiPlus2__dEdx_FDC) ? PiPlus2__dEdx_FDC : 0.","#pi+2 dEdX FDC","",'F');
   //factory->AddVariable("PiPlus2_dEdx_TOF := (PiPlus2__dEdx_TOF == PiPlus2__dEdx_TOF) ? PiPlus2__dEdx_TOF : 0.","#pi+2 dEdX TOF","",'F');
   //factory->AddVariable("PiPlus2_dEdx_ST := (PiPlus2__dEdx_ST == PiPlus2__dEdx_ST) ? PiPlus2__dEdx_ST : 0.","#pi+2 dEdX ST","",'F');
   //factory->AddVariable("PiPlus2_TrackDOCAToCalorimeterShower := (PiPlus2__TrackDOCA_BCAL < PiPlus2__TrackDOCA_FCAL) ? PiPlus2__TrackDOCA_BCAL : PiPlus2__TrackDOCA_FCAL", "Closest calorimeter shower to #pi+2 track", "", 'F');
   //factory->AddVariable("PiPlus2_TotalCalorimeterEnergy := PiPlus2__Energy_BCAL + PiPlus2__Energy_FCAL","#pi+2 total calorimeter energy","",'F');

   factory->AddVariable("Photon1_TrackDOCAToCalorimeterShower := (Photon1__TrackDOCA_BCAL < Photon1__TrackDOCA_FCAL) ? Photon1__TrackDOCA_BCAL : Photon1__TrackDOCA_FCAL", "Closest calorimeter shower to #gamma1 track", "", 'F');
   //factory->AddVariable("Photon1_TotalCalorimeterEnergy := Photon1__Energy_BCAL + Photon1__Energy_FCAL","#gamma1 total calorimeter energy","",'F');

   factory->AddVariable("Photon2_TrackDOCAToCalorimeterShower := (Photon2__TrackDOCA_BCAL < Photon2__TrackDOCA_FCAL) ? Photon2__TrackDOCA_BCAL : Photon2__TrackDOCA_FCAL", "Closest calorimeter shower to #gamma2 track", "", 'F');
   //factory->AddVariable("Photon2_TotalCalorimeterEnergy := Photon2__Energy_BCAL + Photon2__Energy_FCAL","#gamma2 total calorimeter energy","",'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("Pi0_Mass_Measured := MeasuredPi0Mass_MinusPi0Mass", "Measured Pi0 Mass", "", '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'); 
   factory->AddVariable("Unused__TotalEnergy_Calorimeters := (Unused__TotalEnergy_BCAL + Unused__TotalEnergy_FCAL)", "Total unused Calorimeter 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->AddSignalTree    ( background, signalWeight );
   factory->AddBackgroundTree( background_combinatoric, signalWeight );
   factory->AddBackgroundTree( background, backgroundWeight );
   //Float_t pi0Mean = 0.133 , pi0Sigma = 0.012, missingMassMean = -0.041 , missingMassSigma = 0.11 , omegaMassMean = 0.786, omegaMassSigma = 0.032;
   //Float_t nSigma = 10.0;
   TString bothCut = " && PiPlus1__NDF_Tracking>3 && PiPlus2__NDF_Tracking>3 && PiMinus1__NDF_Tracking>3 && PiMinus2__NDF_Tracking>3"; //Form(" && MeasuredPi0Mass_MinusPi0Mass < %f && MeasuredPi0Mass_MinusPi0Mass > %f && Measured__MissingMass < %f && Measured__MissingMass > %f && Omega_MeasuredMass < %f && Omega_MeasuredMass > %f && NumUnused < 20 ", (pi0Mean + nSigma * pi0Sigma), (pi0Mean - nSigma * pi0Sigma), (missingMassMean + nSigma * missingMassSigma), (missingMassMean - nSigma * missingMassSigma), (omegaMassMean + nSigma * omegaMassSigma), (omegaMassMean - nSigma * omegaMassSigma)); //" && 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=0:nTest_Signal=0:nTest_Background=0:SplitMode=Random:!V" );

	   // book MVA method
   factory->BookMethod( TMVA::Types::kBDT, "BDT","!H:!V:NTrees=500:MaxDepth=1:VarTransform=G,D,G,D" );
   factory->BookMethod( TMVA::Types::kBDT, "BDT_2","!H:!V:NTrees=400:nEventsMin=400:MaxDepth=3:BoostType=AdaBoost:SeparationType=GiniIndex:nCuts=20:PruneMethod=NoPruning" );

   //factory->BookMethod( TMVA::Types::kTMlpANN, "TMlp_ANN", "" );
   //factory->BookMethod( TMVA::Types::kBDT, "BDT","!H:!V:NTrees=1000");
   //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;

}