#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 "TMath.h"


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


//root -l -b -q 'trainingOSMuon_test.C("skim_BToJPsiKMu.root")'

void trainingOSMuon_test(
		TString DataFile1 = ""

	){
	
	TFile *file1 = new TFile(DataFile1);

	// TREES
	TTree *tree1 = (TTree*)file1 ->Get("tree");

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

	// OUTPUT
	TString outfileName( "TMVA_OSMuon.root" );
	TFile* outputFile = TFile::Open( outfileName, "RECREATE" );

	// FACTORY
	std::string factoryOptions( "!V:!Silent:Color:DrawProgressBar:AnalysisType=Classification" );
	TMVA::Factory *factory = new TMVA::Factory( "TMVAClassification", outputFile, factoryOptions );

	//DATALOADER
	TMVA::DataLoader *dataloader=new TMVA::DataLoader("dataset");


	dataloader->AddVariable( "muoPtRel", 'F' );
	dataloader->AddVariable( "muoQCone", 'F' );


	dataloader->AddSpectator( "osMuon", 'I' );


	// Adding tree to factory
	dataloader->AddSignalTree( tree1 );

	dataloader->AddBackgroundTree( tree1 );


	// Apply additional cuts on the signal and background samples (can be different)
	
	TCut mycutgen =	"osMuon==1";
	TCut mycuts = 	"sign(ssBLund)== -1*sign(muoCharge)" + mycutgen;
	TCut mycutb = 	"sign(ssBLund)== +1*sign(muoCharge)" + mycutgen;

	// Telling the factory how to use the training and testing events
	//
	// If no numbers of events are given, half of the events in the tree are used 
	// for training, and the other half for testing:
	//    factory->PrepareTrainingAndTestTree( mycut, "SplitMode=random:!V" );

	dataloader->PrepareTrainingAndTestTree( mycuts, mycutb, "SplitMode=Random:V");

	

	//BDT

	TString bdtOptions("H:V:UseBaggedBoost:BaggedSampleFraction=0.6:NTrees=500:MaxDepth=20:nCuts=-1:MinNodeSize=0.2%:"
						"BoostType=RealAdaBoost:AdaBoostBeta=0.6:DoBoostMonitor=True");


	factory->BookMethod( dataloader, TMVA::Types::kBDT, "BDT_OSMuon", bdtOptions );


    //CATEGORIES
	// The variable sets
	TString JetCatVars("muoPtRel:muoQCone");

	TString noJetCatVars("muoQCone");

	// Categorised classifier
	TMVA::MethodCategory* mcat = 0;
	TMVA::MethodBase* BDT_OSMuon_Cat 	= factory->BookMethod( dataloader, TMVA::Types::kCategory, "BDT_OSMuon_Cat","" );
	mcat = dynamic_cast<TMVA::MethodCategory*>(BDT_OSMuon_Cat);
	mcat->AddMethod( "muoPtRel>=0", JetCatVars, TMVA::Types::kBDT, "BDT_OSMuon_Jet", bdtOptions );
	mcat->AddMethod( "muoPtRel<0", noJetCatVars, TMVA::Types::kBDT, "BDT_OSMuon_noJet", bdtOptions );




	//factory->OptimizeAllMethods("ROCIntegral","Minuit");

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

}