/*
 *  MVAAllUpdate.C
 *
 *  Created on: Dec 13, 2018
 *      Author: John Russell
 *
 *  Compared to the other "KL<MVA_method>Update.C" scripts, this consolidates
 *  all of them for comparison when running over the entire data set.
 */


#include "AnitaTMVA.h"  //  Use this to take care of requisite TMVA includes.


void MVAAllUpdate() {

	//  Open file with trees to be considered in TMVA.
	TFile refinedWAISHPolSumFile("fullSampleSummaries/refinedWAISHPol/summary-full-refinedWAISHPol-part_all-sinsub_10_3_ad_2.root");
	TFile refinedWAISVPolSumFile("fullSampleSummaries/refinedWAISVPol/summary-full-refinedWAISVPol-part_all-sinsub_10_3_ad_2.root");
	TFile refinedHiCal2ASumFile("fullSampleSummaries/refinedHiCal2A/summary-full-refinedHiCal2A-part_all-sinsub_10_3_ad_2.root");
	TFile refinedHiCal2BSumFile("fullSampleSummaries/refinedHiCal2B/summary-full-refinedHiCal2B-part_all-sinsub_10_3_ad_2.root");
	TFile thermal0SumFile("fullSampleSummaries/thermal/summary-full-thermal-part_all_0-sinsub_10_3_ad_2.root");
	TFile thermal1SumFile("fullSampleSummaries/thermal/summary-full-thermal-part_all_1-sinsub_10_3_ad_2.root");
	TFile thermal2SumFile("fullSampleSummaries/thermal/summary-full-thermal-part_all_2-sinsub_10_3_ad_2.root");
	TFile thermal3SumFile("fullSampleSummaries/thermal/summary-full-thermal-part_all_3-sinsub_10_3_ad_2.root");
	TFile minBiasSumFile("fullSampleSummaries/minBias/summary-full-minBias-part_all-sinsub_10_3_ad_2.root");
	TFile payloadBlastSumFile("fullSampleSummaries/payloadBlast/summary-full-payloadBlast-part_all-sinsub_10_3_ad_2.root");
	TFile strongCWSumFile("fullSampleSummaries/strongCW/summary-full-strongCW-part_all-sinsub_10_3_ad_2.root");

	TFile refinedWAISHPolKLCoeffsFile("sampleKLCoeffs/refinedWAISHPol/KLCoeffs-refinedWAISHPol-part_all.root");
	TFile refinedWAISVPolKLCoeffsFile("sampleKLCoeffs/refinedWAISVPol/KLCoeffs-refinedWAISVPol-part_all.root");
	TFile refinedHiCal2AKLCoeffsFile("sampleKLCoeffs/refinedHiCal2A/KLCoeffs-refinedHiCal2A-part_all.root");
	TFile refinedHiCal2BKLCoeffsFile("sampleKLCoeffs/refinedHiCal2B/KLCoeffs-refinedHiCal2B-part_all.root");
	TFile thermalKLCoeffsFile("sampleKLCoeffs/thermal/KLCoeffs-thermal-part_all.root");
	TFile minBiasKLCoeffsFile("sampleKLCoeffs/minBias/KLCoeffs-minBias-part_all.root");
	TFile payloadBlastKLCoeffsFile("sampleKLCoeffs/payloadBlast/KLCoeffs-payloadBlast-part_all.root");
	TFile strongCWKLCoeffsFile("sampleKLCoeffs/strongCW/KLCoeffs-strongCW-part_all.root");

	//  Get pointers to trees from files.
	TTree * refinedWAISHPolSumTree = (TTree *) refinedWAISHPolSumFile.Get("refinedWAISHPol");
	TTree * refinedWAISVPolSumTree = (TTree *) refinedWAISVPolSumFile.Get("refinedWAISVPol");
	TTree * refinedHiCal2ASumTree = (TTree *) refinedHiCal2ASumFile.Get("refinedHiCal2A");
	TTree * refinedHiCal2BSumTree = (TTree *) refinedHiCal2BSumFile.Get("refinedHiCal2B");
	TTree * thermal0SumTree = (TTree *) thermal0SumFile.Get("thermal");
	TTree * thermal1SumTree = (TTree *) thermal1SumFile.Get("thermal");
	TTree * thermal2SumTree = (TTree *) thermal2SumFile.Get("thermal");
	TTree * thermal3SumTree = (TTree *) thermal3SumFile.Get("thermal");
	TTree * minBiasSumTree = (TTree *) minBiasSumFile.Get("minBias");
	TTree * payloadBlastSumTree = (TTree *) payloadBlastSumFile.Get("payloadBlast");
	TTree * strongCWSumTree = (TTree *) strongCWSumFile.Get("strongCW");

	TTree * refinedWAISHPolKLCoeffsTree = (TTree *) refinedWAISHPolKLCoeffsFile.Get("refinedWAISHPolKLCoeffsTree");
	TTree * refinedWAISVPolKLCoeffsTree = (TTree *) refinedWAISVPolKLCoeffsFile.Get("refinedWAISVPolKLCoeffsTree");
	TTree * refinedHiCal2AKLCoeffsTree = (TTree *) refinedHiCal2AKLCoeffsFile.Get("refinedHiCal2AKLCoeffsTree");
	TTree * refinedHiCal2BKLCoeffsTree = (TTree *) refinedHiCal2BKLCoeffsFile.Get("refinedHiCal2BKLCoeffsTree");
	TTree * thermalKLCoeffsTree = (TTree *) thermalKLCoeffsFile.Get("thermalKLCoeffsTree");
	TTree * minBiasKLCoeffsTree = (TTree *) minBiasKLCoeffsFile.Get("minBiasKLCoeffsTree");
	TTree * payloadBlastKLCoeffsTree = (TTree *) payloadBlastKLCoeffsFile.Get("payloadBlastKLCoeffsTree");
	TTree * strongCWKLCoeffsTree = (TTree *) strongCWKLCoeffsFile.Get("strongCWKLCoeffsTree");

	int numThermal = thermal0SumTree -> GetEntries() + thermal1SumTree -> GetEntries() + thermal2SumTree -> GetEntries() + thermal3SumTree -> GetEntries();
	double thermalWeight = 1. / numThermal;  // The thermal trees should be uniformly weighted, as described in https://sourceforge.net/p/tmva/mailman/message/35306050/.

	//  Set up DataLoader objects.
	//  For Fisher HPol signal.
	TMVA::DataLoader * hPolIdealFisherDL = new TMVA::DataLoader("AllResults");

	hPolIdealFisherDL -> AddSignalTree(refinedWAISHPolSumTree);
	hPolIdealFisherDL -> AddSignalTree(refinedHiCal2ASumTree);
	hPolIdealFisherDL -> AddSignalTree(refinedHiCal2BSumTree);

	hPolIdealFisherDL -> AddBackgroundTree(thermal0SumTree, thermalWeight);
	hPolIdealFisherDL -> AddBackgroundTree(thermal1SumTree, thermalWeight);
	hPolIdealFisherDL -> AddBackgroundTree(thermal2SumTree, thermalWeight);
	hPolIdealFisherDL -> AddBackgroundTree(thermal3SumTree, thermalWeight);

	hPolIdealFisherDL -> AddSpectator("run");
	hPolIdealFisherDL -> AddSpectator("eventNumber");

	TMVA::DataLoader * hPolFisherDL = hPolIdealFisherDL -> MakeCopy("AllResults");

	hPolFisherDL -> AddBackgroundTree(minBiasSumTree);
	hPolFisherDL -> AddBackgroundTree(payloadBlastSumTree);
	hPolFisherDL -> AddBackgroundTree(strongCWSumTree);
	hPolFisherDL -> AddBackgroundTree(refinedWAISVPolSumTree);

	hPolFisherDL -> AddSpectator("run");  //  Added trees appear to be copied, but not spectator variables!
	hPolFisherDL -> AddSpectator("eventNumber");

	//  For Fisher VPol signal.
	TMVA::DataLoader * vPolIdealFisherDL = new TMVA::DataLoader("AllResults");

	vPolIdealFisherDL -> AddSignalTree(refinedWAISVPolSumTree);

	vPolIdealFisherDL -> AddBackgroundTree(thermal0SumTree, thermalWeight);
	vPolIdealFisherDL -> AddBackgroundTree(thermal1SumTree, thermalWeight);
	vPolIdealFisherDL -> AddBackgroundTree(thermal2SumTree, thermalWeight);
	vPolIdealFisherDL -> AddBackgroundTree(thermal3SumTree, thermalWeight);

	vPolIdealFisherDL -> AddSpectator("run");
	vPolIdealFisherDL -> AddSpectator("eventNumber");

	TMVA::DataLoader * vPolFisherDL = vPolIdealFisherDL -> MakeCopy("AllResults");

	vPolFisherDL -> AddBackgroundTree(minBiasSumTree);
	vPolFisherDL -> AddBackgroundTree(payloadBlastSumTree);
	vPolFisherDL -> AddBackgroundTree(strongCWSumTree);
	vPolFisherDL -> AddBackgroundTree(refinedWAISHPolSumTree);
	vPolFisherDL -> AddBackgroundTree(refinedHiCal2ASumTree);
	vPolFisherDL -> AddBackgroundTree(refinedHiCal2BSumTree);

	vPolFisherDL -> AddSpectator("run");
	vPolFisherDL -> AddSpectator("eventNumber");

	//  For BDT HPol signal.
	TMVA::DataLoader * hPolBDTDL = new TMVA::DataLoader("AllResults");

	hPolBDTDL -> AddSignalTree(refinedWAISHPolKLCoeffsTree);
	hPolBDTDL -> AddSignalTree(refinedHiCal2AKLCoeffsTree);
	hPolBDTDL -> AddSignalTree(refinedHiCal2BKLCoeffsTree);

	hPolBDTDL -> AddBackgroundTree(thermalKLCoeffsTree);
	hPolBDTDL -> AddBackgroundTree(minBiasKLCoeffsTree);
	hPolBDTDL -> AddBackgroundTree(payloadBlastKLCoeffsTree);
	hPolBDTDL -> AddBackgroundTree(strongCWKLCoeffsTree);
	hPolBDTDL -> AddBackgroundTree(refinedWAISVPolKLCoeffsTree);

	hPolBDTDL -> AddSpectator("run");
	hPolBDTDL -> AddSpectator("eventNumber");

	TMVA::DataLoader * hPolCorrBDTDL = hPolBDTDL -> MakeCopy("AllResults");

	hPolCorrBDTDL -> AddSpectator("run");
	hPolCorrBDTDL -> AddSpectator("eventNumber");

	//  For BDT VPol signal.
	TMVA::DataLoader * vPolBDTDL = new TMVA::DataLoader("AllResults");

	vPolBDTDL -> AddSignalTree(refinedWAISVPolKLCoeffsTree);

	vPolBDTDL -> AddBackgroundTree(thermalKLCoeffsTree);
	vPolBDTDL -> AddBackgroundTree(minBiasKLCoeffsTree);
	vPolBDTDL -> AddBackgroundTree(payloadBlastKLCoeffsTree);
	vPolBDTDL -> AddBackgroundTree(strongCWKLCoeffsTree);
	vPolBDTDL -> AddBackgroundTree(refinedWAISHPolKLCoeffsTree);
	vPolBDTDL -> AddBackgroundTree(refinedHiCal2AKLCoeffsTree);
	vPolBDTDL -> AddBackgroundTree(refinedHiCal2BKLCoeffsTree);

	vPolBDTDL -> AddSpectator("run");
	vPolBDTDL -> AddSpectator("eventNumber");

	TMVA::DataLoader * vPolCorrBDTDL = vPolBDTDL -> MakeCopy("AllResults");

	vPolCorrBDTDL -> AddSpectator("run");
	vPolCorrBDTDL -> AddSpectator("eventNumber");

	//  Variables to train on.
	TString impulsivityMeas = "mostImpulsiveDeconvolvedFiltered(2).impulsivityMeasure";
	TString hPolSignalUnblindedKLCoeffs, vPolSignalUnblindedKLCoeffs, hPolSignalUnblindedCorrKLCoeffs, vPolSignalUnblindedCorrKLCoeffs;

	hPolIdealFisherDL -> AddVariable(impulsivityMeas);
	hPolFisherDL -> AddVariable(impulsivityMeas);
	vPolIdealFisherDL -> AddVariable(impulsivityMeas);
	vPolFisherDL -> AddVariable(impulsivityMeas);
	for (int idx = 0; idx < 20; ++idx) {

		hPolSignalUnblindedKLCoeffs.Form("hPolSignalUnblindedKLCoeffs[%d]", idx);
		vPolSignalUnblindedKLCoeffs.Form("vPolSignalUnblindedKLCoeffs[%d]", idx);
		hPolSignalUnblindedCorrKLCoeffs.Form("hPolSignalUnblindedCorrKLCoeffs[%d]", idx);
		vPolSignalUnblindedCorrKLCoeffs.Form("vPolSignalUnblindedCorrKLCoeffs[%d]", idx);

		hPolBDTDL -> AddVariable(hPolSignalUnblindedKLCoeffs);
		vPolBDTDL -> AddVariable(vPolSignalUnblindedKLCoeffs);
		hPolCorrBDTDL -> AddVariable(hPolSignalUnblindedCorrKLCoeffs);
		vPolCorrBDTDL -> AddVariable(vPolSignalUnblindedCorrKLCoeffs);
	}

	//  Prepare DataLoaders for number of background events equaling number of input events.
//	hPolIdealFisherDL -> PrepareTrainingAndTestTree("eventNumber % 5 == 0", "");
//	hPolFisherDL -> PrepareTrainingAndTestTree("eventNumber % 5 == 0", "");
//	vPolIdealFisherDL -> PrepareTrainingAndTestTree("eventNumber % 5 == 0", "");
//	vPolFisherDL -> PrepareTrainingAndTestTree("eventNumber % 5 == 0", "");
//
//	hPolBDTDL -> PrepareTrainingAndTestTree("eventNumber % 5 == 0", "");
//	hPolCorrBDTDL -> PrepareTrainingAndTestTree("eventNumber % 5 == 0", "");
//	vPolBDTDL -> PrepareTrainingAndTestTree("eventNumber % 5 == 0", "");
//	vPolCorrBDTDL -> PrepareTrainingAndTestTree("eventNumber % 5 == 0", "");

	hPolIdealFisherDL -> PrepareTrainingAndTestTree("", 40000, 40000, 40000, 40000);
	hPolFisherDL -> PrepareTrainingAndTestTree("", 40000, 40000, 40000, 40000);
	vPolIdealFisherDL -> PrepareTrainingAndTestTree("", 40000, 40000, 40000, 40000);
	vPolFisherDL -> PrepareTrainingAndTestTree("", 40000, 40000, 40000, 40000);

	hPolBDTDL -> PrepareTrainingAndTestTree("", 40000, 40000, 40000, 40000);
	hPolCorrBDTDL -> PrepareTrainingAndTestTree("", 40000, 40000, 40000, 40000);
	vPolBDTDL -> PrepareTrainingAndTestTree("", 40000, 40000, 40000, 40000);
	vPolCorrBDTDL -> PrepareTrainingAndTestTree("", 40000, 40000, 40000, 40000);

	//  Create files in which to store TMVA results.
	TFile allMVAFile("AllResults/AllMVA40K.root", "recreate");
//	TFile impulsivityFisherFile("AllResults/impulsivityFisher.root", "recreate");
//	TFile KLCoeffsBDTFile("AllResults/KLCoeffsBDT.root", "recreate");

	//  Run TMVA factory, placing output into output file.
	TMVA::Factory * allMVAFactory = new TMVA::Factory("allMVA40KFactory", & allMVAFile, "V");  //  The "V" is for verbose output.
//	TMVA::Factory * impulsivityFisherFactory = new TMVA::Factory("impulsivityFisherFactory", & impulsivityFisherFile, "V");  //  The "V" is for verbose output.
//	TMVA::Factory * KLCoeffsBDTFactory = new TMVA::Factory("KLCoeffsBDTFactory", & KLCoeffsBDTFile, "V");

	// Book methods to perform testing.
	allMVAFactory -> BookMethod(hPolIdealFisherDL, TMVA::Types::kFisher, "hPolIdealFisher", "H");
	allMVAFactory -> BookMethod(vPolIdealFisherDL, TMVA::Types::kFisher, "vPolIdealFisher", "H");
	allMVAFactory -> BookMethod(hPolFisherDL, TMVA::Types::kFisher, "hPolFisher", "H");
	allMVAFactory -> BookMethod(vPolFisherDL, TMVA::Types::kFisher, "vPolFisher", "H");

	allMVAFactory -> BookMethod(hPolBDTDL, TMVA::Types::kBDT, "hPolBDT", "H:DoBoostMonitor=True:MaxDepth=2:nCuts=200");
	allMVAFactory -> BookMethod(vPolBDTDL, TMVA::Types::kBDT, "vPolBDT", "H:DoBoostMonitor=True:MaxDepth=2:nCuts=200");
	allMVAFactory -> BookMethod(hPolCorrBDTDL, TMVA::Types::kBDT, "hPolCorrBDT", "H:DoBoostMonitor=True:MaxDepth=2:nCuts=200");
	allMVAFactory -> BookMethod(vPolCorrBDTDL, TMVA::Types::kBDT, "vPolCorrBDT", "H:DoBoostMonitor=True:MaxDepth=2:nCuts=200");

//	impulsivityFisherFactory -> BookMethod(hPolIdealFisherDL, TMVA::Types::kFisher, "hPolIdealFisher", "H");
//	impulsivityFisherFactory -> BookMethod(vPolIdealFisherDL, TMVA::Types::kFisher, "vPolIdealFisher", "H");
//	impulsivityFisherFactory -> BookMethod(hPolFisherDL, TMVA::Types::kFisher, "hPolFisher", "H");
//	impulsivityFisherFactory -> BookMethod(vPolFisherDL, TMVA::Types::kFisher, "vPolFisher", "H");
//
//	KLCoeffsBDTFactory -> BookMethod(hPolBDTDL, TMVA::Types::kBDT, "hPolBDT", "H:DoBoostMonitor=True");
//	KLCoeffsBDTFactory -> BookMethod(vPolBDTDL, TMVA::Types::kBDT, "vPolBDT", "H:DoBoostMonitor=True");
//	KLCoeffsBDTFactory -> BookMethod(hPolCorrBDTDL, TMVA::Types::kBDT, "hPolCorrBDT", "H:DoBoostMonitor=True");
//	KLCoeffsBDTFactory -> BookMethod(vPolCorrBDTDL, TMVA::Types::kBDT, "vPolCorrBDT", "H:DoBoostMonitor=True");

	//  Collection of training, testing, and evaluating the MVA, then save the results to tree.
	allMVAFactory -> TrainAllMethods();
	allMVAFactory -> TestAllMethods();
	allMVAFactory -> EvaluateAllMethods();

	allMVAFile.cd();
	allMVAFile.Write();
	allMVAFile.Close();

//	impulsivityFisherFactory -> TrainAllMethods();
//	impulsivityFisherFactory -> TestAllMethods();
//	impulsivityFisherFactory -> EvaluateAllMethods();
//
//	impulsivityFisherFile.cd();
//	impulsivityFisherFile.Write();
//	impulsivityFisherFile.Close();
//
//	KLCoeffsBDTFactory -> TrainAllMethods();
//	KLCoeffsBDTFactory -> TestAllMethods();
//	KLCoeffsBDTFactory -> EvaluateAllMethods();
//
//	KLCoeffsBDTFile.cd();
//	KLCoeffsBDTFile.Write();
//	KLCoeffsBDTFile.Close();

	//  Close other files.
	refinedWAISHPolSumFile.Close();
	refinedWAISVPolSumFile.Close();
	refinedHiCal2ASumFile.Close();
	refinedHiCal2BSumFile.Close();
	minBiasSumFile.Close();
	payloadBlastSumFile.Close();
	strongCWSumFile.Close();
	thermal0SumFile.Close();
	thermal1SumFile.Close();
	thermal2SumFile.Close();
	thermal3SumFile.Close();

	refinedWAISHPolKLCoeffsFile.Close();
	refinedWAISVPolKLCoeffsFile.Close();
	refinedHiCal2AKLCoeffsFile.Close();
	refinedHiCal2BKLCoeffsFile.Close();
	minBiasKLCoeffsFile.Close();
	payloadBlastKLCoeffsFile.Close();
	strongCWKLCoeffsFile.Close();
	thermalKLCoeffsFile.Close();
}