/*
V * StupidScript.C
 *
 *  Created on: 17 Feb 2017
 *      Author: ncl
 */

#include <TChain.h>
#include <TString.h>
#include <TFile.h>
#include <TGraph.h>
#include <TH1F.h>
#include <iostream>
#include <algorithm>
#include <string>
#include <fstream>

double GetVirtualMemoryUsage() {
	std::ifstream file;
	file.open("/proc/self/stat");
	TString datum;
	double VirtualMemory = 0.;
	if (file.is_open()) {
		for (int i = 0; i < 23; i++)
			datum.ReadToDelim(file, ' ');
		VirtualMemory = datum.Atoll() / (1024. * 1024.);
	}
	file.close();
	return VirtualMemory;
}

void FillVirtualMemoryUsage(TGraph* memoryGraph, int x, double&val) {
	/// Fill the TGraph at position x with the current usage of virtual memory
	int npoints = memoryGraph->GetN();
	val = GetVirtualMemoryUsage();
	memoryGraph->SetPoint(npoints, x, val);
}

void FillBasketSize(TGraph* memoryGraph, int x, double val) {
	/// Fill the TGraph at position x with the current usage of virtual memory
	int npoints = memoryGraph->GetN();
	memoryGraph->SetPoint(npoints, x, val);
}

void ProcessBranches(TChain *in){
	/// Disable all branches by default, then re-enable only the ones we need
	TObjArray * arr = in->GetListOfBranches();
	for (int i = 0; i < arr->GetEntries(); ++i) {
		TBranch *b = (TBranch*) arr->At(i);
		b->ResetAddress();
		b->SetStatus(0);
	}

	in->GetBranch("Cedar")->SetStatus(1);
}

void processMemoryPlot(TGraph *fVirtualMemoryVsExportEvent, double memoryMax, double jumpLimit){
	/// Compute the derivative of the memory usage to find jumps. Then print the biggest memory increase.
	std::cout << "Maximum memory used:" << memoryMax << std::endl;

	TH1D *deriv = new TH1D("deriv", "deriv", fVirtualMemoryVsExportEvent->GetN(), 0, fVirtualMemoryVsExportEvent->GetN());
	double x1, y1, x2, y2;
	for(int i=1; i<fVirtualMemoryVsExportEvent->GetN()-1; ++i){
		fVirtualMemoryVsExportEvent->GetPoint(i, x1, y1);
		fVirtualMemoryVsExportEvent->GetPoint(i+1, x2, y2);
		deriv->Fill(x1, (y2-y1)/(x2-x1));
	}

	double max = deriv->GetMaximum();
	int maxBin = deriv->GetMaximumBin();
	double x, y_1, y_2;
	fVirtualMemoryVsExportEvent->GetPoint(maxBin-1, x, y_1);
	fVirtualMemoryVsExportEvent->GetPoint(maxBin, x, y_2);
	std::cout << "Event " << maxBin << " jumps of " << max << " from " << y_1 << " to " << y_2 << (max>jumpLimit ? " -> JUMP" : "") << std::endl;
	deriv->Write();
}

void createMemoryPlots(TGraph* fVirtualMemoryVsExportEvent){
	/// Create the memory plots with appropriate names and titles
	fVirtualMemoryVsExportEvent->SetName("fVirtualMemoryVsExportEvent");
	fVirtualMemoryVsExportEvent->SetTitle("Virtual memory vs. Exported event");
}

void step0(TString inFile, TString outFile){
	/// Extracts the first entry of the input ROOT file and copy it in the output file.
	TGraph* fVirtualMemoryVsExportEvent = new TGraph();
	createMemoryPlots(fVirtualMemoryVsExportEvent);

	TChain *in = new TChain("Reco");
	in->AddFile(inFile.Data());

	ProcessBranches(in);

	TFile *out = TFile::Open(outFile, "RECREATE");
	TTree *recoOut = in->CloneTree(0);

	double memoryMax=0;

	in->GetEntry(0);
	recoOut->Fill();
	FillVirtualMemoryUsage(fVirtualMemoryVsExportEvent, 0, memoryMax);

	processMemoryPlot(fVirtualMemoryVsExportEvent, memoryMax, 40);

	recoOut->Write();
	out->Flush();
	out->Purge();
	out->Close();
}

void step1(TString inFile, TString outFile){
	/// Multiply the first event present in the input ROOT file 1000 times to increase statistics in single file for step2.
	TGraph* fVirtualMemoryVsExportEvent = new TGraph();
	createMemoryPlots(fVirtualMemoryVsExportEvent);

	TChain *in = new TChain("Reco");
	in->AddFile(inFile.Data());

	ProcessBranches(in);

	TFile *out = TFile::Open(outFile, "RECREATE");
	TTree *recoOut = in->CloneTree(0);

	double memoryMax=0;

	unsigned int nevt = 1000;
	in->GetEntry(0);
	for(int i=0; i<nevt; ++i){
		if(i%1000==0)
			std::cout << i << "/" << nevt << std::endl;

		recoOut->Fill();
		FillVirtualMemoryUsage(fVirtualMemoryVsExportEvent, i, memoryMax);
	}

	recoOut->Write();
	fVirtualMemoryVsExportEvent->Write();

	processMemoryPlot(fVirtualMemoryVsExportEvent, memoryMax, 40);

	out->Flush();
	out->Purge();
	out->Close();
}

void step2(TString inFile, TString outFile){
	/// Use the input ROOT file 40 times in the input TChain
	/// and go through all events while copying them in the output
	/// TTree.
	TGraph* fVirtualMemoryVsExportEvent = new TGraph();
	createMemoryPlots(fVirtualMemoryVsExportEvent);

	TChain *in = new TChain("Reco");
	for(int i=0; i<40; ++i)
		in->AddFile(inFile.Data());

	ProcessBranches(in);

	TFile *out = TFile::Open(outFile, "RECREATE");
	TTree *recoOut = in->CloneTree(0);

	double memoryMax=0;

	unsigned int nevt = in->GetEntries();
	for(int i=0; i<nevt; ++i){
		if(i%1000==0)
			std::cout << i << "/" << nevt << std::endl;

		in->GetEntry(i);
		recoOut->Fill();
		FillVirtualMemoryUsage(fVirtualMemoryVsExportEvent, i, memoryMax);
	}

	recoOut->Write();
	fVirtualMemoryVsExportEvent->Write();

	processMemoryPlot(fVirtualMemoryVsExportEvent, memoryMax, 40);

	out->Flush();
	out->Purge();
	out->Close();
}

int main(int argc, char** argv){
	/// Use as ./my_script stepNumber inputFile outputFile
	if(argc!=4)
		return 0;
	int step = std::stoi(argv[1]);

	switch(step){
	case 0:
		step0(argv[2], argv[3]);
		break;
	case 1:
		step1(argv[2], argv[3]);
		break;
	case 2:
		step2(argv[2], argv[3]);
		break;
	}
}