/*
 * makeKLCoeffsOriginal.C
 *
 *  Created on: Jan 31, 2020
 *      Author: John Russell
 *
 *  Same idea as "makeKLCoeffs.C", except the waveform constructed from the
 *  most impulsive peak over both interferometric maps is used instead.
 */
 
 
#include "FFTtools.h"
#include "FFTtoolsRev.h"
#include "AnitaDataset.h"
#include "Hical2.h"


//  Structure from which coefficients will be stored, and related quantites.
struct branchStruct {
	
	double KLCoeffs[2][5];  //  First index corresponds to if polarity flipped, second index to principal componennt order.
	double KLMag[2];  //  Magnitude over KLCoeffs, index correspond to polarity flipped.
	unsigned int KLCoeff0Idx;  //  Based on sise of KLCoeffs[][0], 1 when KLCoeffs[1][0] > KLCoeffs[0][0], 0 otherwise.
	unsigned int KLMagIdx;  //  Based on size of KLMag, 1 when KLMag[1] > KLMag[0], 0 otherwise. Should correspond to polarity flip when KLMagIdx == 1.
	
} hPolCoh, hPolDeconv, hPol, vPolCoh, vPolDeconv, vPol;

//  Open KL bases files.
TFile separateHPolFile("~/makeKLBases/KL_Bases/separateKLBasesHPol.root");
TFile hPolFile("~/makeKLBases/KL_Bases/KLBasesHPol.root");
TFile separateVPolFile("~/makeKLBases/KL_Bases/separateKLBasesVPol.root");
TFile vPolFile("~/makeKLBases/KL_Bases/KLBasesVPol.root");

//  Access TTrees within KL basis files.
TTree * hPolCohTree = (TTree *) separateHPolFile.Get("hPolCohKLBasisGraphTree");
TTree * hPolDeconvTree = (TTree *) separateHPolFile.Get("hPolDeconvKLBasisGraphTree");
TTree * hPolTree = (TTree *) hPolFile.Get("hPolKLBasisGraphTree");
TTree * vPolCohTree = (TTree *) separateVPolFile.Get("vPolCohKLBasisGraphTree");
TTree * vPolDeconvTree = (TTree *) separateVPolFile.Get("vPolDeconvKLBasisGraphTree");
TTree * vPolTree = (TTree *) vPolFile.Get("vPolKLBasisGraphTree");


//  For basis using either coherently-summed or deconvolved graph.
void fillSeparateBranch(branchStruct & KLBasisStruct, TTree * KLBasisTree, UCorrelator::Analyzer & analyzer) {

	TString treeName = TString(KLBasisTree -> GetName());
	
	//  Access vector<TGraph> objects from each file.
	vector<TGraph> * separateGraphs = 0;
	
	if (treeName.Contains("Coh")) KLBasisTree -> SetBranchAddress("cohGraphs", & separateGraphs);
	else if (treeName.Contains("Deconv")) KLBasisTree -> SetBranchAddress("deconvGraphs", & separateGraphs);
	else return 1;
	
	int mostImpulsivePolAsInt = analyzer.getSummary() -> mostImpulsivePolAsInt(2);
	int mostImpulsiveInd = analyzer.getSummary() -> mostImpulsiveInd(2);

	//  Set up array of event TGraphs.
	TGraph eventGraphs[2];
	if (treeName.Contains("Coh") && mostImpulsivePolAsInt == 0) {

		eventGraphs[0] = TGraph(* analyzer.getCoherent(AnitaPol::kHorizontal, mostImpulsiveInd) -> even());
		eventGraphs[1] = TGraph(* analyzer.getCoherentXpol(AnitaPol::kHorizontal, mostImpulsiveInd) -> even());
		
	} else if (treeName.Contains("Coh") && mostImpulsivePolAsInt == 1) {

		eventGraphs[0] = TGraph(* analyzer.getCoherentXpol(AnitaPol::kVertical, mostImpulsiveInd) -> even());
		eventGraphs[1] = TGraph(* analyzer.getCoherent(AnitaPol::kVertical, mostImpulsiveInd) -> even());

	} else if (treeName.Contains("Deconv") && mostImpulsivePolAsInt == 0) {
		
		eventGraphs[0] = TGraph(* analyzer.getDeconvolved(AnitaPol::kHorizontal, mostImpulsiveInd) -> even());
		eventGraphs[1] = TGraph(* analyzer.getDeconvolvedXpol(AnitaPol::kHorizontal, mostImpulsiveInd) -> even());
		
	} else if (treeName.Contains("Deconv") && mostImpulsivePolAsInt == 1) {
		
		eventGraphs[0] = TGraph(* analyzer.getDeconvolvedXpol(AnitaPol::kVertical, mostImpulsiveInd) -> even());
		eventGraphs[1] = TGraph(* analyzer.getDeconvolved(AnitaPol::kVertical, mostImpulsiveInd) -> even());
		
	} else return 2;

	//  Get pointers Y values of event graphs.
	double * YEvent[2];
	for (int i = 0; i < 2; ++i) YEvent[i] = eventGraphs[i].GetY();

	//  Total event power needed for normalization.
	double totPow = 0;
	for (int grNum = 0; grNum < 2; ++grNum) {

		for (int j = 0; j < eventGraphs[grNum].GetN(); ++j) totPow += YEvent[grNum][j] * YEvent[grNum][j];
	}

	double norm = sqrt(totPow);

	//  Filling in the branch.
	double maxLagTime, minLagTime;
	double KLMagSq[2] = {};
	for (int i = 0; i < 5; ++i) {
		
		KLBasisTree -> GetEntry(i);
		
		TGraph * hPolCov = FFTtools::getCovGraph(& eventGraphs[0], & separateGraphs -> at(0));
		TGraph * vPolCov = FFTtools::getCovGraph(& eventGraphs[1], & separateGraphs -> at(1));
		
		TGraph cov = TGraph(* hPolCov);
		int NCov = cov.GetN();
		double * YCov = cov.GetY();
		for (int j = 0; j < NCov; ++j) YCov[j] += vPolCov -> GetY()[j];  //  Since cross-covariance/correlation are calculated to next power of 2 in length of the larger two input graphs, assuming resulting length is the same.

		if (i == 0) {
			
			double * XCov = cov.GetX();
				
			maxLagTime = XCov[TMath::LocMax(NCov, YCov)];
			minLagTime = XCov[TMath::LocMin(NCov, YCov)];
		}

		KLBasisStruct.KLCoeffs[0][i] = cov.Eval(maxLagTime, 0, "S") / norm;
		KLBasisStruct.KLCoeffs[1][i] = -cov.Eval(minLagTime, 0, "S") / norm;
				
		KLMagSq[0] += KLBasisStruct.KLCoeffs[0][i] * KLBasisStruct.KLCoeffs[0][i];
		KLMagSq[1] += KLBasisStruct.KLCoeffs[1][i] * KLBasisStruct.KLCoeffs[1][i];
		
		delete hPolCov, delete vPolCov;
	}

	KLBasisStruct.KLMag[0] = sqrt(KLMagSq[0]);
	KLBasisStruct.KLMag[1] = sqrt(KLMagSq[1]);
	KLBasisStruct.KLCoeff0Idx = KLBasisStruct.KLCoeffs[1][0] > KLBasisStruct.KLCoeffs[0][0];
	KLBasisStruct.KLMagIdx = KLBasisStruct.KLMag[1] > KLBasisStruct.KLMag[0];
}


//  For basis using both coherently-summed and deconvolved graphs.
void fillTotalBranch(branchStruct & KLBasisStruct, TTree * KLBasisTree, UCorrelator::Analyzer & analyzer) {
	
	TString treeName = TString(KLBasisTree -> GetName());
	
	//  Access vector<TGraph> objects from each file.
	vector<TGraph> * cohGraphs = 0, * deconvGraphs = 0;
	
	KLBasisTree -> SetBranchAddress("cohGraphs", & cohGraphs);
	KLBasisTree -> SetBranchAddress("deconvGraphs", & deconvGraphs);

	int mostImpulsivePolAsInt = analyzer.getSummary() -> mostImpulsivePolAsInt(2);
	int mostImpulsiveInd = analyzer.getSummary() -> mostImpulsiveInd(2);

	//  Set up array of event TGraphs.
	TGraph eventGraphs[4];
	if (mostImpulsivePolAsInt == 0) {

		eventGraphs[0] = TGraph(* analyzer.getCoherent(AnitaPol::kHorizontal, mostImpulsiveInd) -> even());
		eventGraphs[1] = TGraph(* analyzer.getCoherentXpol(AnitaPol::kHorizontal, mostImpulsiveInd) -> even());
		eventGraphs[2] = TGraph(* analyzer.getDeconvolved(AnitaPol::kHorizontal, mostImpulsiveInd) -> even());
		eventGraphs[3] = TGraph(* analyzer.getDeconvolvedXpol(AnitaPol::kHorizontal, mostImpulsiveInd) -> even());
		
	} else if (mostImpulsivePolAsInt == 1) {

		eventGraphs[0] = TGraph(* analyzer.getCoherentXpol(AnitaPol::kVertical, mostImpulsiveInd) -> even());
		eventGraphs[1] = TGraph(* analyzer.getCoherent(AnitaPol::kVertical, mostImpulsiveInd) -> even());
		eventGraphs[2] = TGraph(* analyzer.getDeconvolvedXpol(AnitaPol::kVertical, mostImpulsiveInd) -> even());
		eventGraphs[3] = TGraph(* analyzer.getDeconvolved(AnitaPol::kVertical, mostImpulsiveInd) -> even());

	} else return 2;

	//  Get pointers Y values of event graphs.
	double * YEvent[4];
	for (int i = 0; i < 4; ++i) YEvent[i] = eventGraphs[i].GetY();

	//  Total event power needed for normalization.
	double cohTotPow = 0, deconvTotPow = 0;
	for (int grNum = 0; grNum < 2; ++grNum) {

		for (int j = 0; j < eventGraphs[grNum].GetN(); ++j) cohTotPow += YEvent[grNum][j] * YEvent[grNum][j];
		for (int j = 0; j < eventGraphs[grNum + 2].GetN(); ++j) deconvTotPow += YEvent[grNum + 2][j] * YEvent[grNum + 2][j];
	}

	double cohNorm = sqrt(cohTotPow);
	double deconvNorm = sqrt(deconvTotPow);
	
	//  Filling in the branch.
	double cohMaxLagTime, cohMinLagTime;
	double deconvMaxLagTime, deconvMinLagTime;
	double KLMagSq[2] = {};
	for (int i = 0; i < 5; ++i) {
		
		KLBasisTree -> GetEntry(i);
		
		TGraph * hPolCohCov = FFTtools::getCovGraph(& eventGraphs[0], & cohGraphs -> at(0));
		TGraph * vPolCohCov = FFTtools::getCovGraph(& eventGraphs[1], & cohGraphs -> at(1));
		TGraph * hPolDeconvCov = FFTtools::getCovGraph(& eventGraphs[2], & deconvGraphs -> at(0));
		TGraph * vPolDeconvCov = FFTtools::getCovGraph(& eventGraphs[3], & deconvGraphs -> at(1));
		
		TGraph cohCov = TGraph(* hPolCohCov);
		int NCohCov = cohCov.GetN();
		double * YCohCov = cohCov.GetY();
		for (int j = 0; j < NCohCov; ++j) YCohCov[j] += vPolCohCov -> GetY()[j];  //  Since cross-covariance/correlation are calculated to next power of 2 in length of the larger two input graphs, assuming resulting length is the same.

		TGraph deconvCov = TGraph(* hPolDeconvCov);
		int NDeconvCov = deconvCov.GetN();
		double * YDeconvCov = deconvCov.GetY();
		for (int j = 0; j < NDeconvCov; ++j) YDeconvCov[j] += vPolDeconvCov -> GetY()[j];

		if (i == 0) {
			
			double * XCohCov = cohCov.GetX();
			
			cohMaxLagTime = XCohCov[TMath::LocMax(NCohCov, YCohCov)];
			cohMinLagTime = XCohCov[TMath::LocMin(NCohCov, YCohCov)];
			
			double * XDeconvCov = deconvCov.GetX();
			
			deconvMaxLagTime = XDeconvCov[TMath::LocMax(NDeconvCov, YDeconvCov)];
			deconvMinLagTime = XDeconvCov[TMath::LocMin(NDeconvCov, YDeconvCov)];
		}

		KLBasisStruct.KLCoeffs[0][i] = cohCov.Eval(cohMaxLagTime, 0, "S") / cohNorm;
		KLBasisStruct.KLCoeffs[0][i] += deconvCov.Eval(deconvMaxLagTime, 0, "S") / deconvNorm;
		KLBasisStruct.KLCoeffs[1][i] = -cohCov.Eval(cohMinLagTime, 0, "S") / cohNorm;
		KLBasisStruct.KLCoeffs[1][i] -= deconvCov.Eval(deconvMinLagTime, 0, "S") / deconvNorm;
				
		KLMagSq[0] += KLBasisStruct.KLCoeffs[0][i] * KLBasisStruct.KLCoeffs[0][i];
		KLMagSq[1] += KLBasisStruct.KLCoeffs[1][i] * KLBasisStruct.KLCoeffs[1][i];
		
		delete hPolCohCov, delete vPolCohCov;
		delete hPolDeconvCov, delete vPolDeconvCov;
	}
	
	KLBasisStruct.KLMag[0] = sqrt(KLMagSq[0]);
	KLBasisStruct.KLMag[1] = sqrt(KLMagSq[1]);
	KLBasisStruct.KLCoeff0Idx = KLBasisStruct.KLCoeffs[1][0] > KLBasisStruct.KLCoeffs[0][0];
	KLBasisStruct.KLMagIdx = KLBasisStruct.KLMag[1] > KLBasisStruct.KLMag[0];
}


void makeKLCoeffsOriginal(const char * sampleSumFileName, const char * createdFileName) {
	
	//  Access summary file from which KL coefficients will be calculated.
	TFile sampleSumFile(sampleSumFileName);
	TTree * sampleSumTree = (TTree *) sampleSumFile.Get("sampleA4");
	AnitaEventSummary * sampleSum = 0;
	sampleSumTree -> SetBranchAddress("summary", & sampleSum);
	int & run = sampleSum -> run;
	unsigned int & eventNumber = sampleSum -> eventNumber;
	int numEntries = sampleSumTree -> GetEntries();

	//  Create/Open TFile of KL coefficients.
	TFile recoverSampleKLCoeffsFile(createdFileName, "update");
	recoverSampleKLCoeffsFile.Write();
	recoverSampleKLCoeffsFile.Close();

	TFile sampleKLCoeffsFile(createdFileName, "update");

	TTree * sampleKLCoeffsTree = 0;

	unsigned int entryCount;	

	const char * structStr = "KLCoeffs[2][5]/D:KLMag[2]:KLCoeff0Idx/i:KLMagIdx";

	if (!sampleKLCoeffsFile.GetNkeys()) {

		sampleKLCoeffsTree = new TTree("KLCoeffsTree", "TTree storing KL coefficients.");
		sampleKLCoeffsTree -> SetAutoSave(numEntries / 100);  //  AutoSave the tree at every 1% completion.

		sampleKLCoeffsTree -> Branch("run", & run);
		sampleKLCoeffsTree -> Branch("eventNumber", & eventNumber);

//		sampleKLCoeffsTree -> AddFriend(sampleSumTree);

		sampleKLCoeffsTree -> Branch("hPolCoh", & hPolCoh, structStr);
		sampleKLCoeffsTree -> Branch("hPolDeconv", & hPolDeconv, structStr);
		sampleKLCoeffsTree -> Branch("hPol", & hPol, structStr);
		
		sampleKLCoeffsTree -> Branch("vPolCoh", & vPolCoh, structStr);
		sampleKLCoeffsTree -> Branch("vPolDeconv", & vPolDeconv, structStr);
		sampleKLCoeffsTree -> Branch("vPol", & vPol, structStr);

		entryCount = 0;

	} else {

		sampleKLCoeffsTree = (TTree *) sampleKLCoeffsFile.Get("KLCoeffsTree");

		sampleKLCoeffsTree -> SetBranchAddress("run", & run);
		sampleKLCoeffsTree -> SetBranchAddress("eventNumber", & eventNumber);

		sampleKLCoeffsTree -> SetBranchAddress("hPolCoh", & hPolCoh);
		sampleKLCoeffsTree -> SetBranchAddress("hPolDeconv", & hPolDeconv);
		sampleKLCoeffsTree -> SetBranchAddress("hPol", & hPol);
		
		sampleKLCoeffsTree -> SetBranchAddress("vPolCoh", & vPolCoh);
		sampleKLCoeffsTree -> SetBranchAddress("vPolDeconv", & vPolDeconv);
		sampleKLCoeffsTree -> SetBranchAddress("vPol", & vPol);

		entryCount = sampleKLCoeffsTree -> GetEntries();
	}

	//  Expedite processing and use sine subtract cache.
	AnitaDataset::DataDirectory dataDirectory = TString(createdFileName).Contains("iceMC", TString::kIgnoreCase) ? AnitaDataset::ANITA_MC_DATA : AnitaDataset::ANITA_ROOT_DATA;
//	if (dataDirectory != AnitaDataset::ANITA_MC_DATA) UCorrelator::SineSubtractFilter::setUseCache(true);

	//  UCorrelator configuration.
	UCorrelator::AnalysisConfig cfg;  //  Default configuration (cfg.deconvolution_method) is with all-pass deconvolution.
	cfg.deconvolution_method = new AnitaResponse::AllPassDeconvolution;  //  Rather than use the band-limited default, use all-pass to account for lack of elephant trunk corrections.
	cfg.response_option = UCorrelator::AnalysisConfig::ResponseA4;  //  The new A4 response.

	//  Set up Analyzer object with filtering.
	UCorrelator::Analyzer analyzer(& cfg, true);  //  "true" is set so "analyzer" doesn't reset each time when drawing.

	FilterStrategy * fStrat = UCorrelator::getStrategyWithKey("sinsub_10_3_ad_2");  //  Sine subtraction filter strategy used throughout UCorrelator.
	fStrat -> addOperation(new UCorrelator::BH13Filter());

	FilterStrategy * fDeconv = new FilterStrategy;  //  Can't be a null pointer.
	fDeconv -> addOperation(new UCorrelator::AntiBH13Filter());  //  Rather than apply BH13Filter to both waveforms and their deconvolutions, then apply AntiBH13 to the deconvolution, its easier to just apply BH13 to just the waveforms.

	analyzer.setExtraFiltersDeconvolved(fDeconv);
	analyzer.setDisallowedAntennas(0, 1ul << 45);  // V-pol antenna 45 is problematic, so we disable it.

	//  Commence filling in the file.
	unsigned int newEntryCount = 0;
	int entryNumStart = 0;
	int unitStep = 1;
	//  Making sure to skip over entries that have already been saved.
	while (newEntryCount < entryCount) {

		++newEntryCount;
		entryNumStart += unitStep;
	}

	for (int entryNum = entryNumStart; entryNum < numEntries; entryNum += unitStep) {

		sampleSumTree -> GetEntry(entryNum);

		int runNum = run;

		//  Set up AnitaDataset object for given run.		
		AnitaDataset d(runNum, false, WaveCalType::kDefault, dataDirectory, AnitaDataset::kRandomizePolarity);

		do {

			//  Use the analyzer declared above to acquire coherently summed waveform for given run and event number.
			d.getEvent(eventNumber);  //  In the run, pointing to the correct eventNumber.

			FilteredAnitaEvent filtEvent(d.useful(), fStrat, d.gps(), d.header());

			AnitaEventSummary eventSum;

			analyzer.analyze(& filtEvent, & eventSum);

			fillSeparateBranch(hPolCoh, hPolCohTree, analyzer);
			fillSeparateBranch(hPolDeconv, hPolDeconvTree, analyzer);
			fillTotalBranch(hPol, hPolTree, analyzer);

			fillSeparateBranch(vPolCoh, vPolCohTree, analyzer);
			fillSeparateBranch(vPolDeconv, vPolDeconvTree, analyzer);
			fillTotalBranch(vPol, vPolTree, analyzer);

			sampleKLCoeffsTree -> Fill();

			//  Further annoyance with the necessity of memory clearance.
			analyzer.clearInteractiveMemory();

			entryNum += unitStep;

			//  Handling end of tree.
			if (entryNum < numEntries) sampleSumTree -> GetEntry(entryNum);
			else break;

		} while (runNum == run);

		entryNum -= unitStep;
	}

	//  Writing to sampleKLCoeffsFile, then closing it.
	sampleKLCoeffsFile.cd();
	sampleKLCoeffsTree -> BuildIndex("run", "eventNumber");
	sampleKLCoeffsTree -> Write();
	sampleKLCoeffsFile.Close();
}