// charged_particle_tree.cc is based on main01.cc from the 
// PYTHIA event generator, but has been modified by
// Peter Christiansen (peter.christiansen@hep.lu.se at Lund University)
//

#include "Pythia8/Pythia.h"

// ROOT includes
#include <TH1.h>
#include <TMath.h>
#include <TFile.h>
#include <TTree.h>
#include <TClonesArray.h>
#include <TParticle.h>

int main() {
  // setup parametes
  const int   nEvents = 1000000;
  const float ymax = 10.0;   

  // Initialize PYTHIA minbias Generator.
  Pythia8::Pythia pythia;
  pythia.readString("Beams:eCM = 7000."); // 7 TeV pp
  /*
    SoftQCD:all = on                   ! Allow total sigma = elastic/SD/DD/ND
    Optionally only study one or a few processes at a time.
    SoftQCD:elastic = on               ! Elastic
    SoftQCD:singleDiffractive = on     ! Single diffractive
    SoftQCD:doubleDiffractive = on     ! Double diffractive
    SoftQCD:centralDiffractive = on    ! Central diffractive
    SoftQCD:nonDiffractive = on        ! Nondiffractive (inelastic)
    SoftQCD:inelastic = on             ! All inelastic
   */
  pythia.readString("SoftQCD:nonDiffractive = on");    
  pythia.init();

  //
  // Create ROOT objects
  //

  // Write histograms to file
  TFile* outFile = new TFile("charged_particle_tree.root", "RECREATE");

  TH1F* hdNdeta = 
    new TH1F("hdNdeta", "dN_{ch}/d#eta; #eta; 1/N_{events} dN_{ch}/d#eta",
	     60, -ymax, ymax);

  TH1F* hMult = 
    new TH1F("hMult", "Total charged particle multiplicity; Mult; Counts",
	     100, 0, 100);

  TTree* tree = new TTree("TT", "PYTHIA Track Tree");
  TClonesArray* trackArray = new TClonesArray("TParticle", 100);
  tree->Bronch("tracks", "TClonesArray", &trackArray);

  int   nRealEvents = 0;
  
  for (int iEvent = 0; iEvent < nEvents; ++iEvent) {
    
    int nCharged  = 0;
    
    if (!pythia.next()) continue;
    
    nRealEvents++;
    
    if (iEvent < 1) {pythia.info.list(); pythia.event.list();} 
    
    for (int i = 0; i < pythia.event.size(); ++i) {
      
      // Final
      if (!pythia.event[i].isFinal())
	continue;
      
      // Hadron
      if(!pythia.event[i].isHadron())
	continue;
      
      // Charged
      if(pythia.event[i].isNeutral())
	continue;

      TParticle* track =
	new((*trackArray)[nCharged]) TParticle();
      track->SetPdgCode(pythia.event[i].id());
      track->SetMomentum(pythia.event[i].px(), pythia.event[i].py(), 
			 pythia.event[i].pz(), pythia.event[i].e());

      nCharged++;

      const float eta = pythia.event[i].eta();
      hdNdeta->Fill(eta);
    }
    
    // Fill charged multiplicity histogram
    hMult->Fill(nCharged);
    // Update tree for this event
    tree->Fill();
    trackArray->Delete();
  } // End of event loop.
  
  // write PYTHIA summary to screen
  pythia.stat();
  
  // Here you could add the code to normalize hdNdeta
  
  
  // Write and close output file
  outFile->Write();
  outFile->Close();
  
  // Check to see that we got most of the events we wanted
  cout << "Real events/simulated: " << nRealEvents 
       << "/ " << nEvents << endl;
  
  return 0;
}