/*
 * getVoltsAtDigitizer.C
 *
 *  Created on: Nov 30, 2019
 *      Author: John Russell
 *
 *  For purposes of simplicity, separating the file production from "getVoltsAtDigitizer.C".
 */


#include "FFTtools.h"
#include "AnitaDataset.h"
#include "AnitaConventions.h"
#include "UCUtil.h"
#include "Hical2.h"


//  Fixing typo.
#ifndef NUM_DIGITIZED_CHANNELS
#define NUM_DIGITIZED_CHANNELS NUM_DIGITZED_CHANNELS
#endif


void writeTreeToFile(const char * dataType, int numRuns, const char * production) {
	
	//  Setting up file in which to save voltage values.
	TFile voltsAtDigitizerFile("voltsAtDigitizer.root", "update");
	
	//  TTree to be placed into file.
	TTree dataTypeTree(TString::Format("%sTree", dataType), TString::Format("Voltage values relative to %s.", dataType));
	
	//  Variables in TTrees of noise RMS file.
	int run;
	dataTypeTree.Branch("run", & run);
	
	unsigned int eventNumber;
	dataTypeTree.Branch("eventNumber", & eventNumber);
	
	double weight;
	dataTypeTree.Branch("weight", & weight);
	
	short tuffIndex;
	dataTypeTree.Branch("tuffIndex", & tuffIndex);
	
	int fNumPoints[NUM_DIGITIZED_CHANNELS];
	dataTypeTree.Branch("fNumPoints", fNumPoints, TString::Format("fNumPoints[%d]/I", NUM_DIGITIZED_CHANNELS));

	double fNoiseAtDigitizer[NUM_DIGITIZED_CHANNELS][NUM_SAMP];
	dataTypeTree.Branch("fNoiseAtDigitizer", fNoiseAtDigitizer, TString::Format("fNoiseAtDigitizer[%d][fNumPoints]/D", NUM_SAMP));

	double fSignalAtDigitizer[NUM_DIGITIZED_CHANNELS][NUM_SAMP];
	dataTypeTree.Branch("fSignalAtDigitizer", fSignalAtDigitizer, TString::Format("fSignalAtDigitizer[%d][fNumPoints]/D", NUM_DIGITIZED_CHANNELS));

	double fVolts[NUM_DIGITIZED_CHANNELS][NUM_SAMP];
	dataTypeTree.Branch("fVolts", fVolts, TString::Format("fVolts[%d][fNumPoints]/D", NUM_DIGITIZED_CHANNELS));

	double fNoiseAtDigitizerRMS[NUM_DIGITIZED_CHANNELS];
	dataTypeTree.Branch("fNoiseAtDigitizerRMS", fNoiseAtDigitizerRMS, TString::Format("fNoiseAtDigitizerRMS[%d]/D", NUM_DIGITIZED_CHANNELS));

	double fSignalAtDigitizerRMS[NUM_DIGITIZED_CHANNELS];
	dataTypeTree.Branch("fSignalAtDigitizerRMS", fSignalAtDigitizerRMS, TString::Format("fSignalAtDigitizerRMS[%d]/D", NUM_DIGITIZED_CHANNELS));

	double fVoltsRMS[NUM_DIGITIZED_CHANNELS];
	dataTypeTree.Branch("fVoltsRMS", fVoltsRMS, TString::Format("fVoltsRMS[%d]/D", NUM_DIGITIZED_CHANNELS));

	double fVoltsRelSignalAtDigitizerRMS[NUM_DIGITIZED_CHANNELS];
	dataTypeTree.Branch("fVoltsRelSignalAtDigitizerRMS", fVoltsRelSignalAtDigitizerRMS, TString::Format("fVoltsRelSignalAtDigitizerRMS[%d]/D", NUM_DIGITIZED_CHANNELS));

	//  Create TChains of the truthAnitaTree for each of the background types.
	TChain truthAnitaChain("truthAnitaTree");
	TChain eventChain("eventTree");

	//  Chaining up the TTrees.
	const char * filePath = "/home/john/Dropbox/UH_Manoa/ANITA/anita_data/anita4_data/iceMC_Kotera_runs/Production";
	for (int runNum = 1; runNum <= numRuns; ++runNum) {

		truthAnitaChain.Add(TString::Format("%s%s/%s/run%d/SimulatedAnitaTruthFile%d.root", filePath, production, dataType, runNum, runNum));
		eventChain.Add(TString::Format("%s%s/%s/run%d/SimulatedAnitaEventFile%d.root", filePath, production, dataType, runNum, runNum));
	}

	//  Variables to be referenced from chains.
	truthAnitaChain.SetMakeClass(1);  //  "[T]o tell the TTree that you want the object to be decomposed into its elementary parts." (https://root-forum.cern.ch/t/ttree-with-multi-dimensional-leaves/11494)
	eventChain.SetMakeClass(1);
	
	truthAnitaChain.SetBranchAddress("run", & run);
	truthAnitaChain.SetBranchAddress("eventNumber", & eventNumber);
	truthAnitaChain.SetBranchAddress("weight", & weight);
	truthAnitaChain.SetBranchAddress("tuffIndex", & tuffIndex);
	truthAnitaChain.SetBranchAddress(TString::Format("fNoiseAtDigitizer[%d][%d]", NUM_DIGITIZED_CHANNELS, NUM_SAMP), fNoiseAtDigitizer);
	truthAnitaChain.SetBranchAddress(TString::Format("fSignalAtDigitizer[%d][%d]", NUM_DIGITIZED_CHANNELS, NUM_SAMP), fSignalAtDigitizer);

	eventChain.SetBranchAddress(TString::Format("fNumPoints[%d]", NUM_DIGITIZED_CHANNELS), fNumPoints);
	eventChain.SetBranchAddress(TString::Format("fVolts[%d][%d]", NUM_DIGITIZED_CHANNELS, NUM_SAMP), fVolts);
	
	//  Filling the created TTrees.
	for (int entryNum = 0; entryNum < truthAnitaChain.GetEntries(); ++entryNum) {
	
		truthAnitaChain.GetEntry(entryNum);
		eventChain.GetEntry(entryNum);

		//  Calculating the RMS over each channel.
		for (int chan = 0; chan < NUM_DIGITIZED_CHANNELS; ++chan){
			
			double fNoiseAtDigitizerSqSum = 0, fSignalAtDigitizerSqSum = 0, fVoltsSqSum = 0, fVoltsRelSignalAtDigitizerSqSum = 0;
			
			for (int samp = 0; samp < fNumPoints[chan]; ++samp) {
			
				fNoiseAtDigitizerSqSum += fNoiseAtDigitizer[chan][samp] * fNoiseAtDigitizer[chan][samp];
				fSignalAtDigitizerSqSum += fSignalAtDigitizer[chan][samp] * fSignalAtDigitizer[chan][samp];
				fVoltsSqSum += fVolts[chan][samp] * fVolts[chan][samp];
				fVoltsRelSignalAtDigitizerSqSum += (fVolts[chan][samp] - fSignalAtDigitizer[chan][samp]) * (fVolts[chan][samp] - fSignalAtDigitizer[chan][samp]);
			}

			fNoiseAtDigitizerRMS[chan] = sqrt(fNoiseAtDigitizerSqSum / fNumPoints[chan]);
			fSignalAtDigitizerRMS[chan] = sqrt(fSignalAtDigitizerSqSum / fNumPoints[chan]);
			fVoltsRMS[chan] = sqrt(fVoltsSqSum / fNumPoints[chan]);
			fVoltsRelSignalAtDigitizerRMS[chan] = sqrt(fVoltsRelSignalAtDigitizerSqSum / fNumPoints[chan]);
		}

		dataTypeTree.Fill();
	}
	
	//  Write the resulting TTrees to file, then close it.
	voltsAtDigitizerFile.cd();	
	dataTypeTree.Write();
	
	//  Now close the written file.
	voltsAtDigitizerFile.Close();
}

void makeVoltsAtDigitizerFile(int numRuns = 10, const char * production = "Capybara") {

	//  Setting up file in which to save voltage values.
	TFile voltsAtDigitizerFile("voltsAtDigitizer.root", "recreate");
	
	//  Now close the written file, as it will be written into using the above function for each data type.
	voltsAtDigitizerFile.Close();
	
	//  Writing the trees to the file.
	writeTreeToFile("Energy_222", numRuns, production);
	writeTreeToFile("MinBiasEnergy_E222", numRuns, production);
	writeTreeToFile("MinBiasEnergyRenorm_E222", numRuns, production);
	writeTreeToFile("MinBiasNoSignalEnergy_E222", numRuns, production);
}