/*
 * root2gamp.cxx
 *
 *  Created on: Jan 17, 2014
 *      Author: yqiang
 */

#include <iostream>
#include <fstream>
using namespace std;

#include <TFile.h>
#include <TROOT.h>
#include <TTree.h>
#include <TBranch.h>
#include <TTreeFormula.h>
#include <TLorentzVector.h>
#include <TString.h>
#include "particleType.h"

/*
 * Structure of particle
 */
struct particle {
	int pid;
	int charge;
};
/*
 * Structure of entry
 */
struct event {
	Long64_t entry;
	double value;
};

char *INPUT = NULL;
char CONFIG[100] = "config.txt";
char OUTPUT[100] = "output.gamp";
char OUTPUTINDEX[100] = "output.gamp.index";
bool APPEND = false;
bool BEAM = false;
bool TARGET = false;
int NMAX = -1;

/*
 * PrintUsage()
 */
int PrintUsage() {
	cerr << "Utility to generate gamp data file from a ROOT tree" << endl;
	cerr << "Author: Yi Qiang <yqiang@jlab.org>" << endl;
	cerr << "Usage: root2gamp [options] <input file>" << endl;
	cerr << "\t-H print help messages" << endl;
	cerr << "\t-B save beam information" << endl;
	cerr << "\t-T save target information" << endl;
	cerr << "\t-A append existing file" << endl;
	cerr << "\t-M<events> max number of events processed (default: all)"
			<< endl;
	cerr << "\t-C<config file> specify config file (default: config.txt)"
			<< endl;
	cerr << "\t-O<output file> specify output filename (default: output.gamp)"
			<< endl;
	return 0;
}

/*
 * ParseArguments
 */
int ParseArguments(int narg, char *argv[]) {
	if (narg < 2) {
		PrintUsage();
		return 1;
	}
	for (int i = 1; i < narg; i++) {
		if (argv[i][0] == '-') {
			char *ptr = &argv[i][1];
			switch (*ptr) {
			case 'h':
			case 'H':
				PrintUsage();
				return 1;
			case 'a':
			case 'A':
				APPEND = true;
				cout << "Appending existing files." << endl;
				break;
			case 'm':
			case 'M':
				sscanf(&ptr[1], "%d", &NMAX);
				cout << "Maximum number of events processed: " << NMAX << endl;
				break;
			case 'c':
			case 'C':
				sscanf(&ptr[1], "%s", CONFIG);
				cout << "Configuration file: " << CONFIG << endl;
				break;
			case 'o':
			case 'O':
				sscanf(&ptr[1], "%s", OUTPUT);
				cout << "Output file name: " << OUTPUT << endl;
				break;
			case 't':
			case 'T':
				TARGET = true;
				cout << "Target information will be saved." << endl;
				break;
			case 'b':
			case 'B':
				BEAM = true;
				cout << "Beam information will be saved." << endl;
				break;
			default:
				cerr << "Unknown option \"" << argv[i] << "\"" << endl;
				PrintUsage();
				return 1;
			}
		} else {
			INPUT = argv[i];
			cout << "Input file: " << INPUT << endl;
		}
	}
	if (!INPUT) {
		PrintUsage();
		return 1;
	}
	return 0;
}

/*
 * Particle ID
 */
particle ParticleID(TString partname) {
	particle part;
	part.charge = 0;
	part.pid = 0;
	if (partname.Contains("Gamma") || partname.Contains("Photon")) {
		part.pid = 1;
	} else if (partname.Contains("Positron")) {
		part.pid = 2;
	} else if (partname.Contains("Electron")) {
		part.pid = 3;
	} else if (partname.Contains("Neutrino")) {
		part.pid = 4;
	} else if (partname.Contains("MuonPlus")) {
		part.pid = 5;
	} else if (partname.Contains("MuonMinus")) {
		part.pid = 6;
	} else if (partname.Contains("Pi0")) {
		part.pid = 7;
	} else if (partname.Contains("PiPlus")) {
		part.pid = 8;
	} else if (partname.Contains("PiMinus")) {
		part.pid = 9;
	} else if (partname.Contains("KLong")) {
		part.pid = 10;
	} else if (partname.Contains("KPlus")) {
		part.pid = 11;
	} else if (partname.Contains("KMinus")) {
		part.pid = 12;
	} else if (partname.Contains("Neutron")) {
		part.pid = 13;
	} else if (partname.Contains("Proton")) {
		part.pid = 14;
	}
	part.charge = ParticleCharge(Particle_t(part.pid));
	return part;
}
/*
 * Main
 */
int main(int narg, char *argv[]) {
	if (ParseArguments(narg, argv))
		return 1;
	/*
	 * Read Configuration
	 */
	TString treename = "";
	TString test_formula = "";
	TString test_condition = "";
	TString beam = "";
	TString target = "";
	int nbranch = 0;
	bool max_formula = true;

	if (!CONFIG) {
		cerr << "No configuration file!" << endl;
		return 1;
	}
	ifstream *cfile = new ifstream(CONFIG);
	if (!cfile->is_open()) {
		cerr << "Unable to open configuration file \"" << CONFIG
				<< "\" for reading." << endl;
		return 1;
	}
	cout << "\nReading config file \"" << CONFIG << "\" ..." << endl;

	(*cfile) >> treename;
	cout << "TTree: " << treename << endl;

	(*cfile) >> test_formula >> test_condition;
	if (test_condition.First("nN") == 2)
		max_formula = false;
	cout << "Formula: " << test_formula << endl;
	cout << "Maximize Formula: ";
	if (max_formula)
		cout << "YES" << endl;
	else
		cout << "NO" << endl;

	(*cfile) >> beam;
	particle beamparticle = ParticleID(beam);
	cout << "Beam: " << beam << "\t ID: " << beamparticle.pid << "\t Charge: "
			<< beamparticle.charge << endl;

	(*cfile) >> target;
	particle targetparticle = ParticleID(target);
	TLorentzVector targetvector = TLorentzVector(0, 0, 0,
			ParticleMass(Particle_t(targetparticle.pid)));
	cout << "Target: " << target << "\t ID: " << targetparticle.pid
			<< "\t Charge: " << targetparticle.charge << endl;

	(*cfile) >> nbranch;
	cout << "Number of Branches: " << nbranch << endl;
	if (nbranch <= 1)
		return 1;
	TString branchname[nbranch];
	particle particles[nbranch];
	for (int i = 0; i < nbranch; i++) {
		(*cfile) >> branchname[i];
		if (branchname[i] == "") {
			cerr << "Missing branch #" << i << endl;
			return 1;
		}
		particles[i] = ParticleID(branchname[i]);
		cout << "Name: " << branchname[i] << "\t ID: " << particles[i].pid
				<< "\t Charge: " << particles[i].charge << endl;
	}
	cfile->close();

	/*
	 * Open ROOT file
	 */
	TFile *rfile = new TFile(INPUT, "r");
	if (!rfile->IsOpen())
		return 1;
	TTree *tree = (TTree*) gROOT->FindObject(treename);
	if (!tree) {
		cerr << "Cannot find tree \"" << treename << "\"" << endl;
		return 1;
	}
	cout << "\nReading ROOT tree \"" << treename << "\" ..." << endl;

	UInt_t evtmax = tree->GetMaximum("EventNumber");
	if (NMAX != -1 && NMAX < evtmax)
		evtmax = UInt_t(NMAX);
	UInt_t evtmin = tree->GetMinimum("EventNumber");
	if (evtmin > evtmax) {
		cerr << "No event within the range specified!" << endl;
		return 1;
	}

	UInt_t eventnumber;
	if (tree->SetBranchAddress("EventNumber", &eventnumber))
		return 1;

	TTreeFormula *tf1 = new TTreeFormula("tf1", test_formula, tree);
	if (!tf1->GetNcodes())
		return 1;

	TLorentzVector *beamvector = new TLorentzVector();
	if (BEAM) {
		if (tree->SetBranchAddress(beam, &beamvector))
			return 1;
	}

	TLorentzVector *vectors[nbranch];
	for (int i = 0; i < nbranch; i++) {
		vectors[i] = new TLorentzVector();
		if (tree->SetBranchAddress(branchname[i], vectors + i))
			return 1;
	}

	tree->SetBranchStatus("*", 0);
	tree->SetBranchStatus("EventNumber", 1);
	for (int i = 0; i < tf1->GetNcodes(); i++) {
		tf1->GetLeaf(i)->GetBranch()->SetStatus(1);
	}

	/*
	 * Go through the tree
	 */
	event events[evtmax];
	for (UInt_t i = 0; i < evtmax; i++) {
		events[i].entry = 0;
		events[i].value = 0;
	}

	Long64_t nentries = tree->GetEntries();
	cout << "Sorting combos based on formula values ..." << endl;
	for (Long64_t i = 0; i < nentries; i++) {
		tree->GetEntry(i);
		double newvalue = tf1->EvalInstance();
//		cout << i << " " << eventnumber << " " << newvalue << endl;
		if (eventnumber > evtmax)
			continue;
		Long64_t oldentry = events[eventnumber - 1].entry;
		double oldvalue = events[eventnumber - 1].value;
		if (oldentry == 0 || (max_formula && newvalue > oldvalue)
				|| (!max_formula && newvalue < oldvalue)) {
			events[eventnumber - 1].value = newvalue;
			events[eventnumber - 1].entry = i + 1;
		}
	}

	tree->SetBranchStatus("*", 0);
	if (BEAM)
		tree->SetBranchStatus(beam, 1);
	for (int i = 0; i < nbranch; i++)
		tree->SetBranchStatus(branchname[i], 1);

	ofstream *gampfile;
	ofstream *indexfile;
	sprintf(OUTPUTINDEX, "%s.index", OUTPUT);

	if (APPEND) {
		gampfile = new ofstream(OUTPUT, ios::app);
		indexfile = new ofstream(OUTPUTINDEX, ios::app);
	} else {
		gampfile = new ofstream(OUTPUT, ios::trunc);
		indexfile = new ofstream(OUTPUTINDEX, ios::trunc);
	}

	int nparticle = nbranch;
	if (BEAM)
		nparticle++;
	if (TARGET)
		nparticle++;

	if (gampfile->is_open() && indexfile->is_open()) {
		cout << "\nWriting GUMP file \"" << OUTPUT << "\" ..." << endl;
		for (UInt_t i = 0; i < evtmax; i++) {
			if (events[i].entry > 0) {
				tree->GetEntry(events[i].entry - 1);
				/*
				 * Write entry to GAMP
				 */
				(*gampfile) << nparticle << "\n";
				if (BEAM) {
					(*gampfile) << beamparticle.pid << " "
							<< beamparticle.charge << " " << beamvector->X()
							<< " " << beamvector->Y() << " " << beamvector->Z()
							<< " " << beamvector->E() << "\n";
				}
				if (TARGET) {
					(*gampfile) << targetparticle.pid << " "
							<< targetparticle.charge << " " << targetvector.X()
							<< " " << targetvector.Y() << " "
							<< targetvector.Z() << " " << targetvector.E()
							<< "\n";
				}
				for (int j = 0; j < nbranch; j++) {
					(*gampfile) << particles[j].pid << " "
							<< particles[j].charge << " " << vectors[j]->X()
							<< " " << vectors[j]->Y() << " " << vectors[j]->Z()
							<< " " << vectors[j]->E() << "\n";
				}
				(*indexfile) << i + 1 << " " << events[i].value << "\n";
			}
			if ((i + 1) % 10000 == 0)
				cout << i + 1 << " (" << double(i + 1) * 100. / evtmax
						<< " %) events processed" << endl;
		}
		gampfile->close();
		indexfile->close();
	}
	return 0;
}