//
// May 12, 2016  D. Lawrence
//
// This macro is used to train a BDT to be used in
// the Level-3 trigger. It assumes there already
// exists a file l3bdt.root that was created using
// the l3bdt plugin like this:
//
//  hd_root -o l3bdt.root -PPLUGINS=L3BDTtree file.evio
//

#include <iostream>

#include<TMVA/Factory.h>

void trainMVA(void)
{
	// Open input file and get tree
	TChain *cppmva = new TChain("cppmva");
	cppmva->AddFile("MUON/hd_root.root");
	cppmva->AddFile("PION/hd_root.root");

	// Open output root file (for TMVA)
	TFile *outfile = new TFile("CPPMVA_out.root", "RECREATE");
	TMVA::Factory *fac = new TMVA::Factory("CPP",outfile,"");

	// Specify input tree that contains both signal and background 
	TCut signalCut("is_pion==1");
	TCut backgroundCut("is_pion==0");
	fac->SetInputTrees(cppmva, signalCut, backgroundCut);

	// Add variables
	fac->AddVariable("Ntracks",              'I');
	fac->AddVariable("Ntof",                 'I');
	fac->AddVariable("Nfcal_hits",           'I');
	fac->AddVariable("Nfcal_clusters",       'I');
	fac->AddVariable("Efcal_clusters",       'F');
	fac->AddVariable("Nfmwpc",               'I');
	fac->AddVariable("Nfmwpc1",              'I');
	fac->AddVariable("Nfmwpc2",              'I');
	fac->AddVariable("Nfmwpc3",              'I');
	fac->AddVariable("Nfmwpc4",              'I');
	fac->AddVariable("Nfmwpc5",              'I');
	fac->AddVariable("Nfmwpc6",              'I');
	fac->AddVariable("Nfmwpc7",              'I');
	fac->AddVariable("Nfmwpc8",              'I');
	fac->AddVariable("Nfmwpc_with_2hits",    'I');
	fac->AddVariable("Nfmwpc_with_lt_2hits", 'I');
	fac->AddVariable("Nfmwpc_with_gt_2hits", 'I');


	TCut preSelectCut("");
	fac->PrepareTrainingAndTestTree(preSelectCut,"");
	fac->BookMethod(TMVA::Types::kBDT, "BDT", "!H:!V:NTrees=850:MinNodeSize=1.72553%:MaxDepth=3:BoostType=AdaBoost:AdaBoostBeta=0.5:SeparationType=GiniIndex:nCuts=20:PruneMethod=NoPruning");
	//fac->BookMethod(TMVA::Types::kKNN, "KNN");

	fac->TrainAllMethods();
	fac->TestAllMethods();
	fac->EvaluateAllMethods();

	delete fac;

	outfile->Close();
	delete outfile;
}