/*
 *  trainKLCoeffsFullResponsesSemiweightedStatic.C
 *
 *  Created on: Mar 18, 2020
 *      Author: John Russell
 *
 *  It appears that running over the full set has too much overhead for
 *  "trainKLCoeffsFullResponsesSemiweihted.C" to run to completion, but
 *  this could be due to certain variables being dynamically allocated
 *  instead of statically. Here, I attempt to change some of these variables
 *  to see if this allows for completion.
 */


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


void trainKLCoeffsFullResponsesSemiweightedStatic() {

	//  Create TChains from which to access eventSummaries for various data types.
	const char * summaryPath = "~/fullSampleFiles/summaries/summary-full-%s-part_all_*-sinsub_10_3_ad_2.root";
	
	TChain waisHPolSumChain("sampleA4"), waisVPolSumChain("sampleA4"), hiCal2ASumChain("sampleA4"), hiCal2BSumChain("sampleA4"), iceMCSumChain("sampleA4");
	TChain thermalSumChain("sampleA4"), minBiasSumChain("sampleA4"), payloadBlastSumChain("sampleA4"), strongCWSumChain("sampleA4");
	
	waisHPolSumChain.Add(TString::Format(summaryPath, "WAISHPol"));
	waisVPolSumChain.Add(TString::Format(summaryPath, "WAISVPol"));
	hiCal2ASumChain.Add(TString::Format(summaryPath, "HiCal2A"));
	hiCal2BSumChain.Add(TString::Format(summaryPath, "HiCal2B"));
	iceMCSumChain.Add(TString::Format(summaryPath, "iceMC-Pangolin-MinBiasEnergy_E222"));

	thermalSumChain.Add(TString::Format(summaryPath, "thermal"));
	minBiasSumChain.Add(TString::Format(summaryPath, "minBias"));
	payloadBlastSumChain.Add(TString::Format(summaryPath, "payloadBlast"));
	strongCWSumChain.Add(TString::Format(summaryPath, "strongCW"));
	
	//  Create TChains from which to access KL coefficients, then friend them with the above corresponding summaries.
	const char * coeffPath = "~/fullSampleFiles/KL_coeffs/KLCoeffsOriginal-full-%s-part_all_*.root";

	TChain waisHPolChain("KLCoeffsTree"), waisVPolChain("KLCoeffsTree"), hiCal2AChain("KLCoeffsTree"), hiCal2BChain("KLCoeffsTree"), iceMCChain("KLCoeffsTree");
	TChain thermalChain("KLCoeffsTree"), minBiasChain("KLCoeffsTree"), payloadBlastChain("KLCoeffsTree"), strongCWChain("KLCoeffsTree");

	waisHPolChain.Add(TString::Format(coeffPath, "WAISHPol"));
	waisVPolChain.Add(TString::Format(coeffPath, "WAISVPol"));
	hiCal2AChain.Add(TString::Format(coeffPath, "HiCal2A"));
	hiCal2BChain.Add(TString::Format(coeffPath, "HiCal2B"));
	iceMCChain.Add(TString::Format(coeffPath, "iceMC-Pangolin-MinBiasEnergy_E222"));

	thermalChain.Add(TString::Format(coeffPath, "thermal"));
	minBiasChain.Add(TString::Format(coeffPath, "minBias"));
	payloadBlastChain.Add(TString::Format(coeffPath, "payloadBlast"));
	strongCWChain.Add(TString::Format(coeffPath, "strongCW"));
	
	waisHPolChain.AddFriend(& waisHPolSumChain);
	waisVPolChain.AddFriend(& waisVPolSumChain);
	hiCal2AChain.AddFriend(& hiCal2ASumChain);
	hiCal2BChain.AddFriend(& hiCal2BSumChain);
	iceMCChain.AddFriend(& iceMCSumChain);
	
	thermalChain.AddFriend(& thermalSumChain);
	minBiasChain.AddFriend(& minBiasSumChain);
	payloadBlastChain.AddFriend(& payloadBlastSumChain);
	strongCWChain.AddFriend(& strongCWSumChain);

	//  Weights to apply to trees. See discussion at https://root-forum.cern.ch/t/with-multiple-signal-and-background-types-handling-one-signal-or-background-type-divided-over-multiple-trees-using-weights/32252/5
	double waisHPolWeight = 1. / waisHPolChain.GetEntries();
	double waisVPolWeight = 1. / waisVPolChain.GetEntries();
	double hiCal2AWeight = 1. / hiCal2AChain.GetEntries();
	double hiCal2BWeight = 1. / hiCal2BChain.GetEntries();
	double iceMCWeight = 1. / iceMCChain.GetEntries();
	
	int totalXPolH = waisHPolChain.GetEntries() + hiCal2AChain.GetEntries() + hiCal2BChain.GetEntries();
	int totalXPolV = waisVPolChain.GetEntries() + iceMCChain.GetEntries();

	//  Set up DataLoader objects.
	//  Hpol objects.
	TMVA::DataLoader hPolIdealDL("trainedResults");

	hPolIdealDL.AddSignalTree(& waisHPolChain, waisHPolWeight);	
	hPolIdealDL.AddSignalTree(& hiCal2AChain, hiCal2AWeight);
	hPolIdealDL.AddSignalTree(& hiCal2BChain, hiCal2BWeight);
	
	hPolIdealDL.AddBackgroundTree(& thermalChain);

	hPolIdealDL.AddSpectator("run");
	hPolIdealDL.AddSpectator("eventNumber");
	hPolIdealDL.AddSpectator("hPolCoh.KLMagIdx");
	hPolIdealDL.AddSpectator("hPolDeconv.KLMagIdx");
	hPolIdealDL.AddSpectator("hPol.KLMagIdx");

	TMVA::DataLoader hPolPracticalDL("trainedResults");
	TMVA::DataLoaderCopy(& hPolPracticalDL, & hPolIdealDL);

	hPolPracticalDL.AddBackgroundTree(& minBiasChain);
	hPolPracticalDL.AddBackgroundTree(& payloadBlastChain);
	hPolPracticalDL.AddBackgroundTree(& strongCWChain);

	hPolPracticalDL.AddSpectator("run");
	hPolPracticalDL.AddSpectator("eventNumber");
	hPolPracticalDL.AddSpectator("hPolCoh.KLMagIdx");
	hPolPracticalDL.AddSpectator("hPolDeconv.KLMagIdx");
	hPolPracticalDL.AddSpectator("hPol.KLMagIdx");
	
	TMVA::DataLoader hPolPracticalXPolDL("trainedResults");
	TMVA::DataLoaderCopy(& hPolPracticalXPolDL, & hPolPracticalDL);

	hPolPracticalXPolDL.AddBackgroundTree(& waisVPolChain, waisVPolWeight * totalXPolV / 2);
	hPolPracticalXPolDL.AddBackgroundTree(& iceMCChain, iceMCWeight * totalXPolV / 2);

	hPolPracticalXPolDL.AddSpectator("run");
	hPolPracticalXPolDL.AddSpectator("eventNumber");
	hPolPracticalXPolDL.AddSpectator("hPolCoh.KLMagIdx");
	hPolPracticalXPolDL.AddSpectator("hPolDeconv.KLMagIdx");
	hPolPracticalXPolDL.AddSpectator("hPol.KLMagIdx");

	//  Vpol objects.
	TMVA::DataLoader vPolIdealDL("trainedResults");
	
	vPolIdealDL.AddSignalTree(& waisVPolChain, waisVPolWeight);
	vPolIdealDL.AddSignalTree(& iceMCChain, iceMCWeight);
	
	vPolIdealDL.AddBackgroundTree(& thermalChain);

	vPolIdealDL.AddSpectator("run");
	vPolIdealDL.AddSpectator("eventNumber");
	vPolIdealDL.AddSpectator("vPolCoh.KLMagIdx");
	vPolIdealDL.AddSpectator("vPolDeconv.KLMagIdx");
	vPolIdealDL.AddSpectator("vPol.KLMagIdx");
	
	TMVA::DataLoader vPolPracticalDL("trainedResults");
	TMVA::DataLoaderCopy(& vPolPracticalDL, & vPolIdealDL);

	vPolPracticalDL.AddBackgroundTree(& minBiasChain);
	vPolPracticalDL.AddBackgroundTree(& payloadBlastChain);
	vPolPracticalDL.AddBackgroundTree(& strongCWChain);

	vPolPracticalDL.AddSpectator("run");
	vPolPracticalDL.AddSpectator("eventNumber");
	vPolPracticalDL.AddSpectator("vPolCoh.KLMagIdx");
	vPolPracticalDL.AddSpectator("vPolDeconv.KLMagIdx");
	vPolPracticalDL.AddSpectator("vPol.KLMagIdx");
	
	TMVA::DataLoader vPolPracticalXPolDL("trainedResults");
	TMVA::DataLoaderCopy(& vPolPracticalXPolDL, & vPolPracticalDL);

	vPolPracticalXPolDL.AddBackgroundTree(& waisHPolChain, waisHPolWeight * totalXPolH / 3);
	vPolPracticalXPolDL.AddBackgroundTree(& hiCal2AChain, hiCal2AWeight * totalXPolH / 3);
	vPolPracticalXPolDL.AddBackgroundTree(& hiCal2BChain, hiCal2BWeight * totalXPolH / 3);

	vPolPracticalXPolDL.AddSpectator("run");
	vPolPracticalXPolDL.AddSpectator("eventNumber");
	vPolPracticalXPolDL.AddSpectator("vPolCoh.KLMagIdx");
	vPolPracticalXPolDL.AddSpectator("vPolDeconv.KLMagIdx");
	vPolPracticalXPolDL.AddSpectator("vPol.KLMagIdx");

	//  Variables to train on.
 	hPolIdealDL.AddVariable("hPolCoh.KLMag[hPolCoh.KLMagIdx]");
 	hPolIdealDL.AddVariable("hPolDeconv.KLMag[hPolDeconv.KLMagIdx]");
 	hPolIdealDL.AddVariable("hPol.KLMag[hPol.KLMagIdx]");
 	hPolIdealDL.AddVariable("mostImpulsiveDeconvolvedFiltered(2).impulsivityMeasure");
 	hPolIdealDL.AddVariable("mostImpulsiveCoherentFiltered(2).impulsivityMeasure");
 	hPolIdealDL.AddVariable("mostImpulsiveDeconvolvedFiltered(2).linearPolFrac()");
  	hPolIdealDL.AddVariable("mostImpulsiveCoherentFiltered(2).linearPolFrac()");

 	hPolPracticalDL.AddVariable("hPolCoh.KLMag[hPolCoh.KLMagIdx]");
 	hPolPracticalDL.AddVariable("hPolDeconv.KLMag[hPolDeconv.KLMagIdx]");
 	hPolPracticalDL.AddVariable("hPol.KLMag[hPol.KLMagIdx]");
 	hPolPracticalDL.AddVariable("mostImpulsiveDeconvolvedFiltered(2).impulsivityMeasure");
 	hPolPracticalDL.AddVariable("mostImpulsiveCoherentFiltered(2).impulsivityMeasure");
 	hPolPracticalDL.AddVariable("mostImpulsiveDeconvolvedFiltered(2).linearPolFrac()");
   	hPolPracticalDL.AddVariable("mostImpulsiveCoherentFiltered(2).linearPolFrac()");
	
 	hPolPracticalXPolDL.AddVariable("hPolCoh.KLMag[hPolCoh.KLMagIdx]");
 	hPolPracticalXPolDL.AddVariable("hPolDeconv.KLMag[hPolDeconv.KLMagIdx]");
 	hPolPracticalXPolDL.AddVariable("hPol.KLMag[hPol.KLMagIdx]");
 	hPolPracticalXPolDL.AddVariable("mostImpulsiveDeconvolvedFiltered(2).impulsivityMeasure");
  	hPolPracticalXPolDL.AddVariable("mostImpulsiveCoherentFiltered(2).impulsivityMeasure");
 	hPolPracticalXPolDL.AddVariable("mostImpulsiveDeconvolvedFiltered(2).linearPolFrac()");
    hPolPracticalXPolDL.AddVariable("mostImpulsiveCoherentFiltered(2).linearPolFrac()");
	
 	vPolIdealDL.AddVariable("vPolCoh.KLMag[vPolCoh.KLMagIdx]");
 	vPolIdealDL.AddVariable("vPolDeconv.KLMag[vPolDeconv.KLMagIdx]");
 	vPolIdealDL.AddVariable("vPol.KLMag[vPol.KLMagIdx]");
 	vPolIdealDL.AddVariable("mostImpulsiveDeconvolvedFiltered(2).impulsivityMeasure");
  	vPolIdealDL.AddVariable("mostImpulsiveCoherentFiltered(2).impulsivityMeasure");
 	vPolIdealDL.AddVariable("mostImpulsiveDeconvolvedFiltered(2).linearPolFrac()");
    vPolIdealDL.AddVariable("mostImpulsiveCoherentFiltered(2).linearPolFrac()");
	
 	vPolPracticalDL.AddVariable("vPolCoh.KLMag[vPolCoh.KLMagIdx]");
 	vPolPracticalDL.AddVariable("vPolDeconv.KLMag[vPolDeconv.KLMagIdx]");
 	vPolPracticalDL.AddVariable("vPol.KLMag[vPol.KLMagIdx]");
 	vPolPracticalDL.AddVariable("mostImpulsiveDeconvolvedFiltered(2).impulsivityMeasure");
  	vPolPracticalDL.AddVariable("mostImpulsiveCoherentFiltered(2).impulsivityMeasure");
 	vPolPracticalDL.AddVariable("mostImpulsiveDeconvolvedFiltered(2).linearPolFrac()");
    vPolPracticalDL.AddVariable("mostImpulsiveCoherentFiltered(2).linearPolFrac()");
	
 	vPolPracticalXPolDL.AddVariable("vPolCoh.KLMag[vPolCoh.KLMagIdx]");
 	vPolPracticalXPolDL.AddVariable("vPolDeconv.KLMag[vPolDeconv.KLMagIdx]");
 	vPolPracticalXPolDL.AddVariable("vPol.KLMag[vPol.KLMagIdx]");
 	vPolPracticalXPolDL.AddVariable("mostImpulsiveDeconvolvedFiltered(2).impulsivityMeasure");
   	vPolPracticalXPolDL.AddVariable("mostImpulsiveCoherentFiltered(2).impulsivityMeasure");
 	vPolPracticalXPolDL.AddVariable("mostImpulsiveDeconvolvedFiltered(2).linearPolFrac()");
    vPolPracticalXPolDL.AddVariable("mostImpulsiveCoherentFiltered(2).linearPolFrac()");

	//  Attempting to cut out NaN values, which should be absent anyway.
	TCut hPolCut("isfinite(hPolCoh.KLMag) && hPolCoh.KLMag <= 1 && isfinite(hPolDeconv.KLMag) && hPolDeconv.KLMag <= 1 && isfinite(hPol.KLMag) && hPol.KLMag <= 2");
	hPolIdealDL.PrepareTrainingAndTestTree(hPolCut, "");
	hPolPracticalDL.PrepareTrainingAndTestTree(hPolCut, "");
	hPolPracticalXPolDL.PrepareTrainingAndTestTree(hPolCut, "");

	TCut vPolCut("isfinite(vPolCoh.KLMag) && vPolCoh.KLMag <= 1 && isfinite(vPolDeconv.KLMag) && vPolDeconv.KLMag <= 1 && isfinite(vPol.KLMag) && vPol.KLMag <= 2");
	vPolIdealDL.PrepareTrainingAndTestTree(vPolCut, "");
	vPolPracticalDL.PrepareTrainingAndTestTree(vPolCut, "");
	vPolPracticalXPolDL.PrepareTrainingAndTestTree(vPolCut, "");
	
	//  Create files in which to store TMVA results.
	TFile hPolFile("trainedResults/hPolKLCoeffsFullSemiweightedResults.root", "recreate");
	TFile vPolFile("trainedResults/vPolKLCoeffsFullSemiweightedResults.root", "recreate");

	//  Run TMVA factory, placing output into output file.
	TMVA::Factory hPolFactory("hPolKLCoeffsFullSemiweightedFactory", & hPolFile);
	TMVA::Factory vPolFactory("vPolKLCoeffsFullSemiweightedFactory", & vPolFile);

	// Hpol booking, training, and writing.
	hPolFactory.BookMethod(& hPolIdealDL, TMVA::Types::kFisher, "ideal");
	hPolFactory.BookMethod(& hPolPracticalDL, TMVA::Types::kFisher, "practical");
	hPolFactory.BookMethod(& hPolPracticalXPolDL, TMVA::Types::kFisher, "practicalXPol");

	hPolFactory.TrainAllMethods();
	hPolFactory.TestAllMethods();
	hPolFactory.EvaluateAllMethods();

	hPolFile.cd();
	hPolFile.Write();
	hPolFile.Close();	

	// Vpol booking, training, and writing.
	vPolFactory.BookMethod(& vPolIdealDL, TMVA::Types::kFisher, "ideal");
	vPolFactory.BookMethod(& vPolPracticalDL, TMVA::Types::kFisher, "practical");
	vPolFactory.BookMethod(& vPolPracticalXPolDL, TMVA::Types::kFisher, "practicalXPol");
	
	vPolFactory.TrainAllMethods();
	vPolFactory.TestAllMethods();
	vPolFactory.EvaluateAllMethods();

	vPolFile.cd();
	vPolFile.Write();
	vPolFile.Close();
}