////////////////////////////////////////////////////////////////////////////////
//
// Macro to pull out AMPT truth information and store it into an
// event-by-event TTree for future analysis.
//
// Heavily stolen from Javier & Bill:
// /phenix+u/jdok/work/offline/analysis/AMPT_Paper/EventPlaneMethodInitialState/FlowSummary.C
// https://github.com/BillTTTTT/Summer2016/tree/master/AMPT_new_version/code_for_all
//
////////////////////////////////////////////////////////////////////////////////
//
// Darren McGlinchey
// 24 Oct 2016
//
// 11/29/2017  includ heavy flavor hadrons   Ming X Liu
//
////////////////////////////////////////////////////////////////////////////////

#include "TMath.h"
#include "TH2.h"
#include "TF1.h"
#include "TStyle.h"
#include "TVector3.h"
#include "TLorentzVector.h"
#include "TTree.h"
#include "TFile.h"
#include "TRandom3.h"

#include <iostream>
#include <fstream>
#include <string>
#include <sstream>
#include <vector>

using namespace std;


//******************************************************************************
// Function headers
//******************************************************************************
void recenter(vector<TVector3> *vp);
float calc_Psi(vector<TVector3> *vp, float n = 2.);
void calcNucleonPlane(vector<TVector3> *proj, vector<TVector3> *targ,
                      vector<float> *ve2, vector<float> *ve3,
                      vector<float> *vpsi2, vector<float> *vpsi3);
int processParton(const char* inFile,
                  vector<float> *vpsi2, vector<float> *vpsi3);
int processNucleon(const char* inFile,
                   vector<float> *ve2, vector<float> *ve3,
                   vector<float> *vpsi2, vector<float> *vpsi3,
                   vector<float> *ipar);
int processHadron(const char* inFile,
                  vector< vector<int> > *pid,
                  vector< vector<TLorentzVector> > *particles);
int writeOscar(const char* outFile,
               vector< vector<int> > *pid,
               vector< vector<TLorentzVector> > *particles,
               vector<TVector3> *vrtx,
               float beamtilt,
               float bcx, float bcy,
               float bsx, float bsy,
               float zmax);
float shift_pt(float pT, float p, TRandom3 *r);
void shift_pt(vector< vector<TLorentzVector> > *vpT, float p, TRandom3 *r);







/**
 * AMPTTruth():
 *
 * param[in] outFile Name of the output root file (will be overwritten)
 * param[in] oscFile Name of the oscar file output
 * param[in] datFile AMPT dat file
 * param[in] parFile AMPT parton information after string melting
 * param[in] geoFile AMPT participant geometry
 * param[in] pshift  Parameter for shifting the pT distribution. 0 is no shift
 * param[in] beamtilt The beamtilt in x vs y
 * param[in] bcx The beam center in x [cm]
 * param[in] bcy The beam center in y [cm]
 * param[in] bsx The beam spread in x [cm]
 * param[in] bsy The beam spread in y [cm]
 * param[in] zmax The maximum zvrtx [-1*zmax, zmax] in [cm]
 */
void AMPTTruth_v226t7(const char* outFile = "ampt_truth.root",
               const char* oscFile = "oscar.dat",
               const char* datFile = "ana/ampt.dat",
               const char* parFile = "ana/parton-initial-afterPropagation.dat",
               const char* geoFile = "ana/npart-xy.dat",
               // const float pshift = 0.0,
               // const float beamtilt = 0.0,
               // const float bcx = 0.0,
               // const float bcy = 0.0,
               // const float bsx = 0.0,
               // const float bsy = 0.0,
               // const float zmax = 0.0)
               const float pshift = 0.0,
               const float beamtilt = -0.001,
               const float bcx = 0.3,
               const float bcy = 0.02,
               const float bsx = 0.03,
               const float bsy = 0.03,
               const float zmax = 15.0)
{

  //==========================================================================//
  // SET RUNNING CONDITIONS
  //==========================================================================//

  //-- print running conditions
  cout << "***********************************************" << endl;
  cout << "****** AMPTTruth():" << endl;
  cout << "** outFile  : " << outFile << endl;
  cout << "** oscFile  : " << oscFile << endl;
  cout << "** datFile  : " << datFile << endl;
  cout << "** parFile  : " << parFile << endl;
  cout << "** geoFile  : " << geoFile << endl;
  cout << "** pshift   : " << pshift << endl;
  cout << "** beamtilt : " << beamtilt << endl;
  cout << "** bcx      : " << bcx << endl;
  cout << "** bcy      : " << bcy << endl;
  cout << "** bsx      : " << bsx << endl;
  cout << "** bsy      : " << bsy << endl;
  cout << "** zmax     : " << zmax << endl;
  cout << "***********************************************" << endl;


  //==========================================================================//
  // DECLARE VARIABLES
  //==========================================================================//

  //-- Vectors to hold information from each file
  vector<float> e2;
  vector<float> e3;
  vector<float> Psi2_nucleon;
  vector<float> Psi3_nucleon;
  vector<float> Psi2_parton;
  vector<float> Psi3_parton;
  vector<float> ipar;
  vector< vector<int> > vpid;
  vector< vector<TLorentzVector> > particles;
  vector<TVector3> vrtx;

  //-- TRandom object for shifting pT
  TRandom3 *rpt = new TRandom3(0);

  //==========================================================================//
  // RETRIEVE INFORMATION FROM EACH FILE AND WRITE THE OSCAR FILE
  //==========================================================================//

  int Nnuc = processNucleon(geoFile, &e2, &e3, &Psi2_nucleon, &Psi3_nucleon, &ipar);
  int Npar = processParton(parFile, &Psi2_parton, &Psi3_parton);
  int Nhad = processHadron(datFile, &vpid, &particles);

  // shift pT if desired
  if ( TMath::Abs(pshift) > 0 )
    shift_pt(&particles, pshift, rpt);

  int Nosc = writeOscar(oscFile, &vpid, &particles, &vrtx,
                        beamtilt,
                        bcx, bcy,
                        bsx, bsy,
                        zmax);

  // check that we processed the same number of events in all cases!
  if (Nnuc != Npar || Nnuc != Nhad || Npar != Nhad || Nnuc != Nosc)
  {
    cout << "ERROR!! AMPT files contain different number of events!" << endl;
    cout << "        " << Nnuc << " in " << geoFile << endl;
    cout << "        " << Npar << " in " << parFile << endl;
    cout << "        " << Nhad << " in " << datFile << endl;
    cout << "        " << Nosc << " in " << oscFile << endl;
    return;
  }

  // triple check that all the relevant vectors have the correct
  // number of entries
  if ( (unsigned int)e2.size() != (unsigned int)Nnuc ||
       (unsigned int)e2.size() != (unsigned int)e3.size() ||
       (unsigned int)e2.size() != (unsigned int)Psi2_nucleon.size() ||
       (unsigned int)e2.size() != (unsigned int)Psi3_nucleon.size() ||
       (unsigned int)e2.size() != (unsigned int)Psi2_parton.size() ||
       (unsigned int)e2.size() != (unsigned int)Psi3_parton.size() ||
       (unsigned int)e2.size() != (unsigned int)vpid.size() ||
       (unsigned int)e2.size() != (unsigned int)particles.size() ||
       (unsigned int)e2.size() != (unsigned int)vrtx.size() ||
       (unsigned int)e2.size() != (unsigned int)ipar.size() )
  {
    cout << "ERROR!! Vectors don't have the same size!!!" << endl;
    cout << "        Nevent       : " << Nnuc << endl;
    cout << "        e2           : " << e2.size() << endl;
    cout << "        e3           : " << e3.size() << endl;
    cout << "        Psi2_nucleon : " << Psi2_nucleon.size() << endl;
    cout << "        Psi3_nucleon : " << Psi3_nucleon.size() << endl;
    cout << "        Psi2_parton  : " << Psi2_parton.size() << endl;
    cout << "        Psi3_parton  : " << Psi3_parton.size() << endl;
    cout << "        vpid         : " << vpid.size() << endl;
    cout << "        particles    : " << particles.size() << endl;
    cout << "        vrtx         : " << vrtx.size() << endl;
    cout << "        ipar         : " << ipar.size() << endl;
    return;
  }

  //==========================================================================//
  // OPEN OUTPUT FILE AND DEFINE TTREE
  //==========================================================================//
  cout << endl;
  cout << "--> Creating output TTree to be written to " << outFile << endl;

  TFile *fout = new TFile(outFile, "RECREATE");
  if (!fout)
  {
    cout << "ERROR!! Unable to create " << outFile << endl;
    return;
  }

  int event;
  float b;
  float vertex[3];
  float ep2;
  float ep3;
  float Psi2n;
  float Psi3n;
  float Psi2p;
  float Psi3p;
  int npart;
  int pid[10000];
  float pt[10000];
  float phi[10000];
  float eta[10000];
  float E[10000];

  fout->cd();
  TTree *t_event = new TTree("t_event", "AMPT event-wise truth");
  t_event->Branch("event", &event, "event/I");
  t_event->Branch("b", &b, "b/F");
  t_event->Branch("vertex", &vertex, "vertex[3]/F");
  t_event->Branch("epsilon2", &ep2, "epsilon2/F");
  t_event->Branch("epsilon3", &ep3, "epsilon3/F");
  t_event->Branch("Psi2_nucleon", &Psi2n, "Psi2_nucleon/F");
  t_event->Branch("Psi3_nucleon", &Psi3n, "Psi3_nucleon/F");
  t_event->Branch("Psi2_parton", &Psi2p, "Psi2_parton/F");
  t_event->Branch("Psi3_parton", &Psi3p, "Psi3_parton/F");
  t_event->Branch("npart", &npart, "npart/I");
  t_event->Branch("pID", &pid, "pID[npart]/I");
  t_event->Branch("pT", &pt, "pT[npart]/F");
  t_event->Branch("eta", &eta, "eta[npart]/F");
  t_event->Branch("phi", &phi, "phi[npart]/F");
  t_event->Branch("E", &E, "E[npart]/F");

  // fill the tree
  for (vector<float>::size_type ievent = 0; ievent < e2.size(); ievent++)
  {
    //-- fill the variables
    event = ievent;
    b = ipar.at(ievent);
    vertex[0] = vrtx.at(ievent).X() * 1e-13;
    vertex[1] = vrtx.at(ievent).Y() * 1e-13;
    vertex[2] = vrtx.at(ievent).Z() * 1e-13;
    ep2 = e2.at(ievent);
    ep3 = e3.at(ievent);
    Psi2n = Psi2_nucleon.at(ievent);
    Psi3n = Psi3_nucleon.at(ievent);
    Psi2p = Psi2_parton.at(ievent);
    Psi3p = Psi3_parton.at(ievent);

    npart = (int)particles.at(ievent).size();
    if ( npart > 10000 )
      {
	cout << "WARNING!! Number of particles greater than array size! npart=" << npart << ". Will only write out 10000." << endl;
	npart = 10000;
      }
    for (int j = 0; j < npart; j++)
    {

      pid[j] = vpid.at(ievent).at(j);
      pt[j]  = particles.at(ievent).at(j).Pt();
      eta[j] = particles.at(ievent).at(j).Eta();
      phi[j] = particles.at(ievent).at(j).Phi();
      E[j]   = particles.at(ievent).at(j).E();
    } // j

    //print information for first few events
    if (ievent < 5 && false)
    {
      cout << endl;
      cout << " ievent: " << ievent << endl;
      cout << " e2: " << ep2 << endl;
      cout << " e3: " << ep3 << endl;
      cout << " Psi2_nucleon: " << Psi2n << endl;
      cout << " Psi3_nucleon: " << Psi3n << endl;
      cout << " Psi2_parton: " << Psi2p << endl;
      cout << " Psi3_parton: " << Psi3p << endl;
      cout << " npart: " << npart << endl;

      int max = (npart < 5) ? npart : 5;

      cout << " pID: ";
      for (int j = 0; j < max; j++)
        cout << pid[j] << " ";
      cout << endl;

      cout << " pT: ";
      for (int j = 0; j < max; j++)
        cout << pt[j] << " ";
      cout << endl;

      cout << " eta: ";
      for (int j = 0; j < max; j++)
        cout << eta[j] << " ";
      cout << endl;

      cout << " phi: ";
      for (int j = 0; j < max; j++)
        cout << phi[j] << " ";
      cout << endl;

      cout << " E: ";
      for (int j = 0; j < max; j++)
        cout << E[j] << " ";
      cout << endl;

    }



    t_event->Fill();
  } // ievent

  //-- Print confirmation
  cout << endl;
  cout << "--> Parsed information from "
       << t_event->GetEntries() << " events "
       << "... writing to " << outFile
       << endl;

  //-- close the file and cleanup
  t_event->Write();
  fout->Close();

}




/**
 * Recenter the spatial coordinates to the center of mass
 */
void recenter(vector<TVector3> *vp)
{
  //-- First we need to find the cm and recenter
  float x_cm = 0.0;
  float y_cm = 0.0;

  for (unsigned int i = 0; i < (unsigned int)vp->size(); i++)
  {
    x_cm += vp->at(i).X();
    y_cm += vp->at(i).Y();
  }

  x_cm = x_cm / (float)vp->size();
  y_cm = y_cm / (float)vp->size();

  for (unsigned int i = 0; i < (unsigned int)vp->size(); i++)
  {
    vp->at(i).SetX(vp->at(i).X() - x_cm);
    vp->at(i).SetY(vp->at(i).Y() - y_cm);
  }

}

/**
 * Calculate the Psi_n plane for a given set of particles
 */
float calc_Psi(vector<TVector3> *vp, float n)
{
  //-- recenter
  recenter(vp);

  //-- Now calculate the Q vectors
  float Qx = 0.0;
  float Qy = 0.0;

  for (unsigned int i = 0; i < (unsigned int)vp->size(); i++)
  {
    float x = vp->at(i).X();
    float y = vp->at(i).Y();
    float r2 = x * x + y * y;
    float phi = TMath::ATan2(y, x);

    Qx += r2 * TMath::Cos(n * phi);
    Qy += r2 * TMath::Sin(n * phi);
  }

  //-- Calculate Psi
  float Psi = (TMath::ATan2(Qy, Qx) + TMath::Pi()) / n;

  return Psi;
}

/**
 * calculate nucleon information, including:
 * epsilon_2, epsilon_3
 * Psi_2, Psi_3
 */
void calcNucleonPlane(vector<TVector3> *proj, vector<TVector3> *targ,
                      vector<float> *ve2, vector<float> *ve3,
                      vector<float> *vpsi2, vector<float> *vpsi3)
{
  bool WithProj = true;

  float x_cm = 0;
  float y_cm = 0;

  //Compute centroid
  for (unsigned int j = 0; j < targ->size(); j++)
  {
    x_cm += targ->at(j).X();
    y_cm += targ->at(j).Y();
  }

  float Count = (float) targ->size();

  if (WithProj) {
    for (unsigned int j = 0; j < proj->size(); j++)
    {
      x_cm += proj->at(j).X();
      y_cm += proj->at(j).Y();
    }
    Count += (float) proj->size();
  }

  x_cm = x_cm / Count;
  y_cm = y_cm / Count;

  //Shift origin to centroid
  for (unsigned int j = 0; j < targ->size(); j++)
  {
    targ->at(j).SetX(targ->at(j).x() - x_cm);
    targ->at(j).SetY(targ->at(j).Y() - y_cm);
  }

  if (WithProj) {
    //Shift origin to centroid
    for (unsigned int j = 0; j < proj->size(); j++)
    {
      proj->at(j).SetX(proj->at(j).X() - x_cm);
      proj->at(j).SetY(proj->at(j).Y() - y_cm);
    }
  }

  //Compute eccentricity
  //JLN - Option to do with blurred Gaussian distributions... (does not change above x_cm,y_cm calculation)
  bool NucleonSmear = true;
  float SmearSigma = 0.4; // units of fm
  TF1 *fradius = new TF1("fradius", "x*TMath::Exp(-x*x/(2*[0]*[0]))", 0.0, 2.0);
  fradius->SetParameter(0, SmearSigma);
  TF1 *fphi = new TF1("fphi", "1.0", 0.0, 2.0 * TMath::Pi());
  int SmearSamplings = 100;
  if (NucleonSmear) Count = 0;

  float qx_1 = 0;
  float qy_1 = 0;
  float qx_2 = 0;
  float qy_2 = 0;
  float qx_3 = 0;
  float qy_3 = 0;
  float phi = 0;
  float rsq = 0;
  float r = 0;

  float e2 = 0;
  float e3 = 0;

  float psi_1 = 0.0;
  float psi_2 = 0.0;
  float psi_3 = 0.0;

  //Loop over all particles in the target nucleus and compute eccentricity
  for (unsigned int i = 0; i < targ->size(); i++)
  {
    if (!NucleonSmear) {
      r = TMath::Sqrt(TMath::Power(targ->at(i).X(), 2) +
                      TMath::Power(targ->at(i).Y(), 2));
      phi = TMath::ATan2(targ->at(i).Y(), targ->at(i).X());

      qx_1 += r * r * TMath::Cos(phi);
      qy_1 += r * r * TMath::Sin(phi);
      qx_2 += r * r * TMath::Cos(2 * phi);
      qy_2 += r * r * TMath::Sin(2 * phi);
      qx_3 += r * r * TMath::Cos(3 * phi);
      qy_3 += r * r * TMath::Sin(3 * phi);
      rsq += r * r;
    }
    if (NucleonSmear) {
      for (int is = 0; is < SmearSamplings; is++) {
        float rtemp   = fradius->GetRandom();
        float phitemp = fphi->GetRandom();
        float xtemp   = targ->at(i).X() + rtemp * TMath::Sin(phitemp);
        float ytemp   = targ->at(i).Y() + rtemp * TMath::Cos(phitemp);

        r    = TMath::Sqrt(xtemp * xtemp + ytemp * ytemp);
        phi  = TMath::ATan2(ytemp, xtemp);

        qx_1 += r * r * TMath::Cos(phi);
        qy_1 += r * r * TMath::Sin(phi);
        qx_2 += r * r * TMath::Cos(2 * phi);
        qy_2 += r * r * TMath::Sin(2 * phi);
        qx_3 += r * r * TMath::Cos(3 * phi);
        qy_3 += r * r * TMath::Sin(3 * phi);
        rsq += r * r;

        Count = Count + 1;  // noting this is a float
      }
    }
  }

  if (WithProj) {
    for (unsigned int i = 0; i < proj->size(); i++)
    {
      if (!NucleonSmear)
      {
        r = TMath::Sqrt(TMath::Power(proj->at(i).X(), 2) +
                        TMath::Power(proj->at(i).Y(), 2));
        phi = TMath::ATan2(proj->at(i).Y(), proj->at(i).X());

        qx_1 += r * r * TMath::Cos(phi);
        qy_1 += r * r * TMath::Sin(phi);
        qx_2 += r * r * TMath::Cos(2 * phi);
        qy_2 += r * r * TMath::Sin(2 * phi);
        qx_3 += r * r * TMath::Cos(3 * phi);
        qy_3 += r * r * TMath::Sin(3 * phi);
        rsq += r * r;
      }
      if (NucleonSmear)
      {
        for (int is = 0; is < SmearSamplings; is++)
        {
          float rtemp   = fradius->GetRandom();
          float phitemp = fphi->GetRandom();
          float xtemp   = proj->at(i).X() + rtemp * TMath::Sin(phitemp);
          float ytemp   = proj->at(i).Y() + rtemp * TMath::Cos(phitemp);

          r    = TMath::Sqrt(xtemp * xtemp + ytemp * ytemp);
          phi  = TMath::ATan2(ytemp, xtemp);

          qx_1 += r * r * TMath::Cos(phi);
          qy_1 += r * r * TMath::Sin(phi);
          qx_2 += r * r * TMath::Cos(2 * phi);
          qy_2 += r * r * TMath::Sin(2 * phi);
          qx_3 += r * r * TMath::Cos(3 * phi);
          qy_3 += r * r * TMath::Sin(3 * phi);
          rsq += r * r;

          Count = Count + 1;  // noting this is a float
        }
      }

    }
  }

  qx_1 = qx_1 / Count;
  qy_1 = qy_1 / Count;
  qx_2 = qx_2 / Count;
  qy_2 = qy_2 / Count;
  qx_3 = qx_3 / Count;
  qy_3 = qy_3 / Count;
  rsq  = rsq / Count;

  e2 = TMath::Sqrt(qx_2 * qx_2 + qy_2 * qy_2) / rsq;
  e3 = TMath::Sqrt(qx_3 * qx_3 + qy_3 * qy_3) / rsq;

  // psi_1 = TMath::ATan2(qy_1, qx_1) + TMath::Pi();
  psi_2 = (TMath::ATan2(qy_2, qx_2) + TMath::Pi()) / 2;
  psi_3 = (TMath::ATan2(qy_3, qx_3) + TMath::Pi()) / 3;

  if (Count == 1) e2 = 0.0;

  ve2->push_back(e2);
  ve3->push_back(e3);
  vpsi2->push_back(psi_2);
  vpsi3->push_back(psi_3);

}





/**
 * process the parton information.
 *
 * The inFile should have the following format:
 *  Event Header:
 *    For isoft=1 (default AMPT):
 *       event number, event-iteration flag, number of partons;
 *    For isoft=4 (String Melting):
 *        event number, event-iteration flag, number of partons,
 *        number of formed baryons, number of formed mesons,
 *        total number of initial particles,
 *        total number of initial particles that cannot enter ZPC.
 *  Parton information:
 *        PYTHIA particle ID number, three-momentum(Px,Py,Pz), mass, and
 *        space-time coordinates(x,y,z,t) of one final particle at freeze-out.
 *
 * param[in] inFile AMPT parton information file
 * param[in, out] psi2 pointer to vector for Psi_2 for each event
 * param[in, out] psi3 pointer to vector for Psi_3 for each event
 * return Number of events processed
 */
int processParton(const char* inFile,
                  vector<float> *vpsi2, vector<float> *vpsi3)
{
  cout << endl;
  cout << "--> Reading parton information from " << inFile << endl;

  ifstream pfile;
  pfile.open(inFile);
  if (!pfile)
  {
    cout << "ERROR!! Can not open " << inFile << endl;
    return 0;
  }

  int nevent = 0;

  //-- event info
  int ievent;
  int npartons;
  int junk;

  vector<TVector3> vpart;

  //-- parton info
  int pid;
  float px, py, pz;
  float mass;
  float x, y, z, t;
  float j; // unused

  while (pfile)
  {


    //Read the event header for partons
    pfile >> ievent >> junk >> npartons >> junk >> junk >> junk >> junk;

    //Avoid float reading the last entry
    if (!pfile) break;


    //Loop over each particle in the event
    for (int ipart = 0; ipart < npartons; ipart++)
    {


      pfile >> pid
            >> px >> py >> pz
            >> mass
            >> x >> y >> z >> t
	    >> j >> j >> j;

      // Cut on the formation time, t < 0.5 fm/c
      if ( t > 0.5 )
        continue;

      float E =  TMath::Sqrt(TMath::Power(px, 2) +
                             TMath::Power(py, 2) +
                             TMath::Power(pz, 2) +
                             TMath::Power(mass, 2));

      TLorentzVector ev(px, py, pz, E);

      if (ev.Pt() < 0.001)
        continue;

      // eta cut
      if ( ev.Eta() > 3)
        continue;

      vpart.push_back(TVector3(x, y, z));

    } // ipart

    vpsi2->push_back(calc_Psi(&vpart, 2));
    vpsi3->push_back(calc_Psi(&vpart, 3));
    nevent++;
    vpart.clear();

  } // while (pfile)

  return nevent;
}


/**
 * process the geometry information
 *
 * File format for each event:
 *    <event number> <iteration flag> <atomic mass proj> <atomic mass targ> <impact parameter>
 *    x y <target/projectile nucleon> <elastic/inelastic coll> <impact paramter> <original flavor ID> <new flavor ID>
 *    ...
 * Note: Collision marker 3 = inelastic
 *
 * param[in] inFile AMPT nucleon information file
 * param[in, out] e2 pointer to vector for epsilon_2 for each event
 * param[in, out] e3 pointer to vector for epsilon_3 for each event
 * param[in, out] psi2 pointer to vector for Psi_2 for each event
 * param[in, out] psi3 pointer to vector for Psi_3 for each event
 * return Number of events processed
 */
int processNucleon(const char* inFile,
                   vector<float> *ve2, vector<float> *ve3,
                   vector<float> *vpsi2, vector<float> *vpsi3,
                   vector<float> *ipar)
{
  cout << endl;
  cout << "--> Reading parton information from " << inFile << endl;

  ifstream nfile;
  nfile.open(inFile);
  if (!nfile)
  {
    cout << "ERROR!! Can not open " << inFile << endl;
    return 0;
  }

  string line;
  int nevent = 0;

  //Header information
  int eventnumber = 0;
  float iterflag = 0;
  int aproj = 0;
  int atarg = 0;
  float impactparameter = 0;

  //Nucleon information
  float x = 0;
  float y = 0;
  int parentnucleus = 0;
  int collcategory = 0;
  float nucleonimpactparameter = 0;
  int initialflavor = 0;
  int finalflavor = 0;

  //Nucleon vectors
  vector<TVector3> particles_proj;
  vector<TVector3> particles_targ;

  while (getline(nfile, line))
  {
    istringstream iss(line);

    //If a new event is found...
    if (iss >> eventnumber >> iterflag >> aproj >> atarg >> impactparameter)
    {
      //This file lists all nucleons in the collision (spectator + wounded)
      for (int i = 1; i <= aproj + atarg; i++)
      {
        getline(nfile, line);
        istringstream iss(line);
        iss >> x >> y >> parentnucleus >> collcategory >> nucleonimpactparameter >> initialflavor >> finalflavor;

        if (collcategory > 0)
        {
          if (parentnucleus > 0)
          {
            particles_proj.push_back(TVector3(x, y, 0));
          }
          else if (parentnucleus < 0)
          {
            particles_targ.push_back(TVector3(x, y, 0));
          }
        }
      } // i

      //Process the particles in the event to compute epsilon2,3
      calcNucleonPlane(&particles_proj, &particles_targ,
                       ve2, ve3, vpsi2, vpsi3);
      ipar->push_back(impactparameter);

      //Reset after event
      nevent++;
      particles_targ.clear();
      particles_proj.clear();

    }
  } // while (getline(nfile, line))

  return nevent;
}


/**
 * process the final state hadron information
 *
 * The inFile should have the following format:
 *   For each event, the first line gives:
 *      event number, test number(=1), number of particles in the event,
 *      impact parameter, total number of participant nucleons in projectile,
 *      total number of participant nucleons in target, number of participant
 *      nucleons in projectile due to elastic collisions, number of
 *      participant nucleons in projectile due to inelastic collisions,
 *      and corresponding numbers in target.
 *      Note that participant nucleon numbers include nucleons participating
 *      in both elastic and inelastic collisions.
 *   Each of the following lines gives:
 *      PYTHIA particle ID number, three-momentum(Px,Py,Pz), mass, and
 *      space-time coordinates(x,y,z,t) of one final particle at freeze-out.
 *
 * param[in] inFile AMPT final state hadron information
 * param[in, out] particles vector containing vector of particles for each event
 * return Number of events processed
 */
int processHadron(const char* inFile,
                  vector< vector<int> > *pid,
                  vector< vector<TLorentzVector> > *particles)
{
  cout << endl;
  cout << "--> Reading particle information from " << inFile << endl;

  ifstream dfile;
  dfile.open(inFile);
  if (!dfile)
  {
    cout << "ERROR!! Can not open " << inFile << endl;
    return 0;
  }

  int nevent = 0;

  // event header
  int evtnumber;
  int testnum;
  int nlist; //Number of particles in output list
  float impactpar;
  int npartproj;
  int nparttarg;
  int npartprojelas;
  int npartprojinelas;
  int nparttargelas;
  int nparttarginelas;
  float junk;

  // particle information
  int partid;
  float pvec[3];
  float mass;
  float spacetime[4];

  vector<int> vpid;
  vector<TLorentzVector> vpart;

  while (dfile)
  {
    //Read the event header
    dfile >> evtnumber >> testnum >> nlist >> impactpar
          >> npartproj >> nparttarg
          >> npartprojelas >> npartprojinelas
          >> nparttargelas >> nparttarginelas
          >> junk;

    //Avoid float reading the last entry
    if (!dfile) break;


    //Loop over each particle in the event
    for (int i = 0; i < nlist; i++)
    {

      dfile >> partid
            >> pvec[0] >> pvec[1] >> pvec[2]
            >> mass
            >> spacetime[0] >> spacetime[1] >> spacetime[2] >> spacetime[3];

      float energy =  sqrt (pow(pvec[0], 2) + pow(pvec[1], 2) + pow(pvec[2], 2) + pow(mass, 2));

      TLorentzVector ev(pvec[0], pvec[1], pvec[2], energy);

      if (ev.Pt() < 0.001)
      {
        continue;
      }

      int KF = TMath::Abs(partid);
      // updated to inlcude heavy flavor hadrons -- MXL 11/29/2017
      if (!(KF == 211 || KF == 321 || KF == 2212 || (KF > 400 && KF < 600) || (KF>4000 && KF < 6000) )) //Pions, kaons and protons, and D, B hadrons
      {
        continue;
      }

      if (TMath::Abs(ev.Eta()) > 5)
        continue;

      vpart.push_back(ev);
      vpid.push_back(partid);

    } // End loop over particles

    // end event
    pid->push_back(vpid);
    particles->push_back(vpart);
    vpart.clear();
    vpid.clear();
    nevent++;
  }

  return nevent;
}


/**
 * Write out the final state particle information into oscar format
 * required as PISA input
 *
 * param[in] outFile Oscar file output name
 * param[in] pid Vector containing vector of particle pid values for each event
 * param[in] particles Vec containing vec of particle info for each event
 * param[in, out] vrtx (x, y, z) vertex information for each event (filled)
 * param[in] beamtilt The beamtilt in x vs y
 * param[in] bcx The beam center in x [cm]
 * param[in] bcy The beam center in y [cm]
 * param[in] bsx The beam spread in x [cm]
 * param[in] bsy The beam spread in y [cm]
 * param[in] zmax The maximum zvrtx [-1*zmax, zmax] in [cm]
 * return The number of events written
 */
int writeOscar(const char* outFile,
               vector< vector<int> > *pid,
               vector< vector<TLorentzVector> > *particles,
               vector<TVector3> *vrtx,
               float beamtilt,
               float bcx, float bcy,
               float bsx, float bsy,
               float zmax)
{
  cout << endl;
  cout << "--> Writing final state particles to oscar file " << outFile << endl;

  ofstream of(outFile);
  if (!of)
  {
    cout << "ERROR!! Unable to write to " << outFile << endl;
    return 0;
  }

  int nevent = 0;

  TRandom3 *r = new TRandom3(0);

  //write the header
  of << "# OSC1999A" << endl;
  of << "# final_id_p_x" << endl;

  //-- Loop over events
  for (vector<int>::size_type ievent = 0; ievent < pid->size(); ievent++)
  {
    // get the number of particles and write the event header
    unsigned int npart = pid->at(ievent).size();

    of << "0 " << npart << endl;

    //-- Get the random event vertex
    // Beam angle: 0.001 in x vs z
    // BC: (0.3, 0.02) all energies
    // BS: (0.03) 200 GeV
    // BS: (0.05)  62 GeV
    // BS: (0.07)  39 GeV
    // BS: (0.10)  20 GeV
    // z : flat
    // don't forget to convert from cm to fm (1e13)
    // float beamtilt = -1 * 0.001;
    float xvrtx, yvrtx, zvrtx; // cm

    // xvrtx = r->Gaus( 0.3, 0.03) * 1e13;
    // yvrtx = r->Gaus(0.02, 0.03) * 1e13;
    // zvrtx = r->Uniform(-15.0, 15.0) * 1e13;
    if (bsx > 0)
      xvrtx = r->Gaus(bcx, bsx) * 1e13;
    else
      xvrtx = bcx * 1e13;
    if (bsy > 0)
      yvrtx = r->Gaus(bcy, bsy) * 1e13;
    else
      yvrtx = bcy * 1e13;
    if (zmax > 0)
      zvrtx = r->Uniform(-1 * zmax, zmax) * 1e13;
    else
      zvrtx = 0.0;

    vrtx->push_back(TVector3(xvrtx, yvrtx, zvrtx));

    for (unsigned int ipart = 0; ipart < npart; ipart++)
    {

      // rotate tracks (this will be propogated back to the truth tree)
      float px = particles->at(ievent).at(ipart).Px();
      float pz = particles->at(ievent).at(ipart).Pz();
      px = px * TMath::Cos(-1 * beamtilt) + pz * TMath::Sin(-1 * beamtilt);
      pz = -1 * px * TMath::Sin(-1 * beamtilt) + pz * TMath::Cos(-1 * beamtilt);
      particles->at(ievent).at(ipart).SetPx(px);
      particles->at(ievent).at(ipart).SetPz(pz);

      TLorentzVector ev = particles->at(ievent).at(ipart);

      of << ipart << " "
         << pid->at(ievent).at(ipart) << " "
         << " 0 "
         << ev.Px() << " "
         << ev.Py() << " "
         << ev.Pz() << " "
         << ev.E() << " "
         << ev.M() << " "
         << xvrtx << " "
         << yvrtx << " "
         << zvrtx << " "
         << " 0 " // other variables
         << endl;
    }

    of << "0 0" << endl;
    nevent++;

  } // ievent

  return nevent;

}

/**
 * Randomly shift the pT of a given particle
 *
 * Shift according to a random pull from a Gaussian with
 * mean = pT * (1 + p)
 * sig  = 0.1
 *
 * param[in] pT Input pT
 * param[in] p  Shift parameter (see above)
 * param[in] r  Pointer to TRandom3 object
 * return randomly shifted pT value
 */
float shift_pt(float pT, float p, TRandom3 *r)
{
  return r->Gaus(pT * (1 + p), 0.1);
}

/**
 * Randomly shift the pT of a vector of particles
 *
 * param[in, out] vpT Vector of pT values. Will be changed
 * param[in] p  Shift parameter (see above)
 * param[in] r  Pointer to TRandom3 object
 */
void shift_pt(vector< vector<TLorentzVector> > *vpT, float p, TRandom3 *r)
{
  if (!r)
  {
    cout << "ERROR!! shift_pt(): r is not a valid pointer" << endl;
    return;
  }

  for (vector<int>::size_type ievent = 0; ievent < vpT->size(); ievent++)
  {
    for (vector<int>::size_type ipart = 0; ipart < vpT->at(ievent).size(); ipart++)
    {

      float pt = shift_pt(vpT->at(ievent).at(ipart).Pt(), p, r);
      float eta = vpT->at(ievent).at(ipart).Eta();
      float phi = vpT->at(ievent).at(ipart).Phi();
      float E = vpT->at(ievent).at(ipart).E();
      vpT->at(ievent).at(ipart).SetPtEtaPhiE(pt, eta, phi, E);

    } // ipart
  } // ievent

  return;
}