#include <cstdlib>
#include <iostream>
#include <map>
#include <string>

#include "TMath.h"
#include "TChain.h"
#include "TFile.h"
#include "TTree.h"
#include "TString.h"
#include "TObjString.h"
#include "TSystem.h"
#include "TROOT.h"
//#include "Functions.h"
#include "TMath.h"

#include "TMVA/Factory.h"
//#include "TMVA/DataLoader.h"
#include "TMVA/Tools.h"
#include "TMVA/TMVAGui.h"

//#include <TLatex.h>

void Training_Z_SP()
{
// ---------------------------------------------------------------------------------------------------
// This loads the library
TMVA::Tools::Instance();

// Default MVA methods to be trained + tested
std::map<std::string,int> Use;

// 1-dimensional likelihood ("naive Bayes estimator")
   Use["Likelihood"]      = 1;
   Use["LikelihoodD"]     = 1; // the "D" extension indicates decorrelated input variables (see option strings)
   Use["LikelihoodPCA"]   = 1; // the "PCA" extension indicates PCA-transformed input variables (see option strings)
   Use["LikelihoodKDE"]   = 1;
   Use["LikelihoodMIX"]   = 1;
//
// Boosted Decision Trees
   Use["BDT"]             = 0; // uses Adaptive Boost
   Use["BDTG"]            = 0; // uses Gradient Boost
   Use["BDTB"]            = 0; // uses Bagging
   Use["BDTD"]            = 0; // decorrelation + Adaptive Boost
   Use["BDTF"]            = 0; // allow usage of fisher discriminant for node splitting
//
// ---------------------------------------------------------------------------------------------------

// Create a ROOT output file where TMVA will store ntuples, histograms, etc.
TFile* outputFile = TFile::Open( "Z_tmvaOUT_SP.root", "RECREATE" );
//TFile* outputFile = TFile::Open( "Z_tmvaOUT_SP_Pol.root", "RECREATE" );

TMVA::Factory *factory = new TMVA::Factory( "MVAnalysis", outputFile,"!V:AnalysisType=Classification");
//TMVA::Factory *factory = new TMVA::Factory( "MVAnalysis", outputFile,"!V:!Silent:Color:DrawProgressBar:Transformations=I;D;P;G,D:AnalysisType=Classification");


TFile *inputS = TFile::Open("/home/saumyen/HEP_Packages/Delphes-3.4.2/ROOT_Files/mono-Z/New_L_4000/Z_SP_M100_L4000_extract.root");
TFile *inputB = TFile::Open("/home/saumyen/HEP_Packages/Delphes-3.4.2/ROOT_Files/mono-Z/New_L_4000/Z_BG_ll_extract.root");

// global event weights per tree (see below for setting event-wise weights)
Double_t signalWeight     = 0.008811;
Double_t backgroundWeight = 0.278254649;

factory->AddSignalTree( (TTree*)inputS->Get("event_tree"), signalWeight );
factory->AddBackgroundTree ( (TTree*)inputB->Get("event_tree"), backgroundWeight );


//factory->AddVariable( "lepPt","#P_{T}^{l}", "GeV",'F');  //Reconstructed Z PT
//factory->AddVariable( "ETAlep","#eta_{l}", "",'F');
//factory->AddVariable( "PHIlep","#phi_{l}", "rad",'F');
factory->AddVariable( "lepM","M(l^{-},l^{+})", "GeV",'F');
factory->AddVariable( "dPHI","d#phi(l^{-}, l^{+})", "rad",'F');
factory->AddVariable( "dRll","dR(l^{-}, l^{+})", "rad",'F');
factory->AddVariable( "axialMET","axialMET", "GeV",'F');
factory->AddVariable( "fracPT","fracPT","",'F');
factory->AddVariable( "MissingET","#slash{E}_{T}", "GeV",'F');
//factory->AddVariable( "ETAmiss","#slash{#eta}", "",'F');
//factory->AddVariable( "PHImiss","#slash{#phi}", "rad",'F');


TCut preselectionCut = "!TMath::IsNaN(axialMET) && !TMath::IsNaN(fracPT) && !TMath::IsNaN(dPHI) && !TMath::IsNaN(ETAlep) && !TMath::IsNaN(MissingET) ";//&& dRll < 2.5 && MissingET > 80 && fracPT > -0.5 && lepM < 200 && abs(dPHI) < 2.3 && abs(axialMET) < 200";

factory->PrepareTrainingAndTestTree( preselectionCut, "nTrain_Signal=3000:nTrain_Background=30000:SplitMode=Random:NormMode=NumEvents:!V" );


// ************ Likelihood ("naive Bayes estimator") ****************

   if (Use["Likelihood"])
      factory->BookMethod( TMVA::Types::kLikelihood, "Likelihood",
                           "H:!V:TransformOutput:PDFInterpol=Spline2:NSmoothSig[0]=20:NSmoothBkg[0]=20:NSmoothBkg[1]=10:NSmooth=1:NAvEvtPerBin=50" );

   // Decorrelated likelihood
   if (Use["LikelihoodD"])
      factory->BookMethod( TMVA::Types::kLikelihood, "LikelihoodD",
                           "!H:!V:TransformOutput:PDFInterpol=Spline2:NSmoothSig[0]=20:NSmoothBkg[0]=20:NSmooth=5:NAvEvtPerBin=50:VarTransform=Decorrelate" );

   // PCA-transformed likelihood
   if (Use["LikelihoodPCA"])
      factory->BookMethod( TMVA::Types::kLikelihood, "LikelihoodPCA",
                           "!H:!V:!TransformOutput:PDFInterpol=Spline2:NSmoothSig[0]=20:NSmoothBkg[0]=20:NSmooth=5:NAvEvtPerBin=50:VarTransform=PCA" );

   // Use a kernel density estimator to approximate the PDFs
   if (Use["LikelihoodKDE"])
      factory->BookMethod( TMVA::Types::kLikelihood, "LikelihoodKDE",
                           "!H:!V:!TransformOutput:PDFInterpol=KDE:KDEtype=Gauss:KDEiter=Adaptive:KDEFineFactor=0.3:KDEborder=None:NAvEvtPerBin=50" );

   // Use a variable-dependent mix of splines and kernel density estimator
   if (Use["LikelihoodMIX"])
      factory->BookMethod( TMVA::Types::kLikelihood, "LikelihoodMIX",
                           "!H:!V:!TransformOutput:PDFInterpolSig[0]=KDE:PDFInterpolBkg[0]=KDE:PDFInterpolSig[1]=KDE:PDFInterpolBkg[1]=KDE:PDFInterpolSig[2]=Spline2:PDFInterpolBkg[2]=Spline2:PDFInterpolSig[3]=Spline2:PDFInterpolBkg[3]=Spline2:KDEtype=Gauss:KDEiter=Nonadaptive:KDEborder=None:NAvEvtPerBin=50" );


// ************ Boosted Decision Trees *************

   if (Use["BDTG"]) // Gradient Boost
      factory->BookMethod( TMVA::Types::kBDT, "BDTG",
                           "!H:!V:NTrees=20:MinNodeSize=5%:BoostType=Grad:Shrinkage=0.10:UseBaggedBoost:BaggedSampleFraction=0.5:nCuts=20:MaxDepth=2" );

   if (Use["BDT"])  // Adaptive Boost
      factory->BookMethod( TMVA::Types::kBDT, "BDT",
                           "!H:!V:NTrees=20:MinNodeSize=5%:MaxDepth=3:BoostType=AdaBoost:AdaBoostBeta=0.5:UseBaggedBoost:BaggedSampleFraction=0.5:SeparationType=GiniIndex:nCuts=20:SigToBkgFraction=1.0:PruneMethod=NoPruning" );

   if (Use["BDTB"]) // Bagging
      factory->BookMethod( TMVA::Types::kBDT, "BDTB",
                           "!H:!V:NTrees=20:BoostType=Bagging:SeparationType=GiniIndex:nCuts=20" );

   if (Use["BDTD"]) // Decorrelation + Adaptive Boost
      factory->BookMethod( TMVA::Types::kBDT, "BDTD",
                           "!H:!V:NTrees=20:MinNodeSize=5%:MaxDepth=3:BoostType=AdaBoost:SeparationType=GiniIndex:nCuts=20:VarTransform=Decorrelate" );

   if (Use["BDTF"])  // Allow Using Fisher discriminant in node splitting for (strong) linearly correlated variables
      factory->BookMethod( TMVA::Types::kBDT, "BDTF",
                           "!H:!V:NTrees=20:MinNodeSize=5%:UseFisherCuts:MaxDepth=3:BoostType=AdaBoost:AdaBoostBeta=0.5:SeparationType=GiniIndex:nCuts=20" );



// Now you can tell the factory to train, test, and evaluate the MVAs
factory->TrainAllMethods();
factory->TestAllMethods();
factory->EvaluateAllMethods();

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

std::cout << "==> Wrote root file: " << outputFile->GetName() << std::endl;
std::cout << "==> TMVAClassification is done!" << std::endl;

delete factory;
// Launch the GUI for the root macros
//if (!gROOT->IsBatch()) TMVA::TMVAGui( outfileName );

//return 0;
}