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

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

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

int TMVAClassification_ggH()
{
	// The explicit loading of the shared libTMVA is done in TMVAlogon.C, defined in .rootrc
	// if you use your private .rootrc, or run from a different directory, please copy the
	// corresponding lines from .rootrc

	// Methods to be processed can be given as an argument; use format:
	//
	//     mylinux~> root -l TMVAClassification.C\(\"myMethod1,myMethod2,myMethod3\"\)

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

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

	std::cout << std::endl;
	std::cout << "==> Start TMVAClassification" << std::endl;

	// --------------------------------------------------------------------------------------------------

	// Here the preparation phase begins

	// Read training and test data
	TString fname;
	fname = "/home/shankha/Packages/root/tutorials/tmva/ggH_root_files/Int_kZZ_OFFSHELL_dghZZp05kZZp368572.root";
	TFile *inputS = TFile::Open( fname );   
	TTree *signalTree       = (TTree*)inputS->Get("Dists");   

	fname = "/home/shankha/Packages/root/tutorials/tmva/ggH_root_files/ggH_4l_SM_OFFSHELL.root";
	TFile *inputB1  = TFile::Open( fname );
	TTree *background1   = (TTree*)inputB1->Get("Dists");

	fname = "/home/shankha/Packages/root/tutorials/tmva/ggH_root_files/qq_4l_SM_POWHEG.root";
	TFile *inputB2  = TFile::Open( fname );
	TTree *background2   = (TTree*)inputB2->Get("Dists");
	
	fname = "/home/shankha/Packages/root/tutorials/tmva/ggH_root_files/lljj_fB_SM.root";
	TFile *inputB3  = TFile::Open( fname );
	TTree *background3   = (TTree*)inputB3->Get("Dists");
	
	fname = "/home/shankha/Packages/root/tutorials/tmva/ggH_root_files/gg_2l2lp_MCFM.root";
	TFile *inputB4  = TFile::Open( fname );
	TTree *background4   = (TTree*)inputB4->Get("Dists");

	fname = "/home/shankha/Packages/root/tutorials/tmva/ggH_root_files/gg_4l_MCFM.root";
	TFile *inputB5  = TFile::Open( fname );
	TTree *background5   = (TTree*)inputB5->Get("Dists");		

	// Create a ROOT output file where TMVA will store ntuples, histograms, etc.
	TString outfileName( "TMVA.root" );
	TFile* outputFile = TFile::Open( outfileName, "RECREATE" );
//	TMVA::Factory *factory = new TMVA::Factory( "TMVAClassification", outputFile,
//			"!V:!Silent:Color:DrawProgressBar:Transformations=I;D;P;G,D:AnalysisType=Classification" );
	TMVA::Factory *factory = new TMVA::Factory( "TMVAClassification", outputFile,
			"!V:!Silent:Color:DrawProgressBar:Transformations=I" );			
	TMVA::DataLoader *dataloader=new TMVA::DataLoader("dataset");


	dataloader->AddVariable("phiZZm",        "phiZZm",        "",'F');
	dataloader->AddVariable("theta1",        "theta1",        "",'F');
	dataloader->AddVariable("theta2",        "theta2",        "",'F');

        dataloader->AddSpectator("weight",       "weight",        "",'F');

	// global event weights per tree (see below for setting event-wise weights)
	Double_t signalWeight      = 1.0;           // XS = -0.19333 fb,                  43274 events surviving;  K-factor = 2.25997
	Double_t backgroundWeight1 = 0.168616453;   // XS = 0.352234 fb (ggH),            47210 events surviving;  K-factor = 2.25997
	Double_t backgroundWeight2 = 5.667382865;   // XS = 0.729328365 fb (qq -> 4l),    1012 events surviving;   K-factor = 0.788032
	Double_t backgroundWeight3 = 2.659406549;   // XS = 0.165117 fb (lljj fake),      565 events surviving;    K-factor = 0.91
	Double_t backgroundWeight4 = 0.051930399;   // XS = 0.00185533 fb (gg -> 2l 2lp), 581 events surviving;    K-factor = 1.62621
	Double_t backgroundWeight5 = 0.025830762;   // XS = 0.00267487 fb (gg -> 4l),     1684 events surviving;   K-factor = 1.62621
	
// Sig = 131.076, Bkg = 458.160283652, 	

        // Total signal cross-section entering BDT = 0.003175751 pb ~ 953 events at 300 ifb.
        // Total background cross-section entering BDT = 3.077950023 pb ~ 923385 events at 300 ifb.

//--- ==================================================================================================
//--- Classifier   (  #signal, #backgr.)  Optimal-cut  S/sqrt(S+B)      NSig      NBkg   EffSig   EffBkg
//--- --------------------------------------------------------------------------------------------------
//---        BDT:  (      953,   923385)       0.3194      9.23704  396.9975  1450.186   0.4166 0.001571
//--- --------------------------------------------------------------------------------------------------

	// You can add an arbitrary number of signal or background trees
	dataloader->AddSignalTree    ( signalTree,     signalWeight );
	dataloader->AddBackgroundTree( background1, backgroundWeight1 );
	dataloader->AddBackgroundTree( background2, backgroundWeight2 );
	dataloader->AddBackgroundTree( background3, backgroundWeight3 );
	dataloader->AddBackgroundTree( background4, backgroundWeight4 );
	dataloader->AddBackgroundTree( background5, backgroundWeight5 );
	
	dataloader->SetSignalWeightExpression    ("weight");	
	
	// Apply additional cuts on the signal and background samples (can be different)
	TCut mycuts = ""; // for example: TCut mycuts = "abs(var1)<0.5 && abs(var2-0.5)<1";
	TCut mycutb = ""; // for example: TCut mycutb = "abs(var1)<0.5";

//	dataloader->PrepareTrainingAndTestTree( mycuts, mycutb,
//			"nTrain_Signal=1000:nTrain_Background=1000:SplitMode=Random:NormMode=NumEvents:!V" );

	//factory->BookMethod(dataloader, TMVA::Types::kBDT, "BDT",
	//"!H:!V:NTrees=8000:MinNodeSize=2.5%:MaxDepth=1:BoostType=AdaBoost:AdaBoostBeta=0.5:UseBaggedBoost:BaggedSampleFraction=0.5:SeparationType=GiniIndex:nCuts=20" );

//	factory->BookMethod(dataloader, TMVA::Types::kBDT, "BDTG",
//		    "!H:!V:NTrees=100:MinNodeSize=2.5%:MaxDepth=1:BoostType=Grad:UseBaggedBoost:BaggedSampleFraction=0.5:SeparationType=GiniIndex:nCuts=20:NegWeightTreatment=Pray" );

	factory->BookMethod(dataloader, TMVA::Types::kBDT, "BDTG",
		    "!H:!V:NTrees=500:MinNodeSize=2.5%:MaxDepth=1:BoostType=Grad:UseBaggedBoost:BaggedSampleFraction=0.5:SeparationType=GiniIndex:nCuts=20:NegWeightTreatment=Pray" );

	// 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();

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

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

}

int main()
{
	return TMVAClassification_ggH();
}