// main00.cc
// Modified from the main01.cc
// which is a part of the PYTHIA event generator.
// Copyright (C) 2008 Torbjorn Sjostrand.
// PYTHIA is licenced under the GNU GPL version 2, see COPYING for details.
// Please respect the MCnet Guidelines, see GUIDELINES for details.

// This is a simple test program.

#include "Pythia.h"
#include "Basics.h"

#include "TROOT.h"
#include "TFile.h"
#include "TH1.h"

bool isBHadron(int id) {
  // This snippet is meant to capture all B hadrons
  // as given in the PDG.
  if (id<0) id*=-1;
  if (id<500) return false;
  return (fmod(id/100,5.)==0.0 || id/1000==5);
}

using namespace Pythia8;

double myRapidity(Vec4& p) {
  return 0.5*log(p.pPlus()/p.pMinus());
}

int main() {
  // Initialize root
  TROOT root("Manuel's ROOT Session","PYTHIA Histograms");

  // Generator. Process selection. LHC initialization. Histogram.
  Pythia pythia;
 
  // Shorthand for some public members of pythia (also static ones).
  //Event& event = pythia.event;
  ParticleDataTable& pdt = pythia.particleData;
  // The cmnd file below contains
  // the Pythia Tune parameters
  // the processes that are turned on
  // and the PDFs used
  // for the pythia run.
 
  pythia.readFile("main00.cmnd");
  UserHooks *oniumUserHook = new SuppressSmallPT();
  pythia.setUserHooksPtr(oniumUserHook);

  cout << "Mass of b quark " << ParticleDataTable::mass(5) << endl;
  cout << "Mass of b bar   " << ParticleDataTable::mass(-5) << endl;
 
  // Extract settings to be used in the main program.
  int    nEvent  = pythia.mode("Main:numberOfEvents");
  int    nList   = pythia.mode("Main:numberToList");
  int    nShow   = pythia.mode("Main:timesToShow");
  int nAllowErr  = pythia.mode("Main:timesAllowErrors");
  bool   showCS  = pythia.flag("Main:showChangedSettings");
  bool showSett  = pythia.flag("Main:showAllSettings");
  bool showStat  = pythia.flag("Main:showAllStatistics");
  bool   showCPD = pythia.flag("Main:showChangedParticleData");
 

  pythia.init();
  if (showSett) pythia.settings.listAll();
  if (showCS) pythia.settings.listChanged();
  if (showCPD) pdt.listChanged();

  Hist mult("charged multiplicity", 100, -0.5, 799.5);
 
  TH1D* multHist = new TH1D("multHist","Multiplicity",100,-0.5,99.5);
  TH1D* bquarkPt = new TH1D("bquarkPt","bquarkPt",100,0,50);
  TH1D* bbarquarkPt = new TH1D("bbarquarkPt","bbar quark Pt",100,0,50);
  TH1D* B0mesonPt = new TH1D("BOmesonPt","B0mesonPt",100,0,50);
  TH1D* B0barmesonPt = new TH1D("BObarmesonPt","B0bar meson Pt",100,0,50);
  TH1D* BplusmesonPt = new TH1D("BplusmesonPt","BplusmesonPt",100,0,50);
  TH1D* BminusmesonPt = new TH1D("BminusmesonPt","Bminus meson Pt",100,0,50);
  TH1D* BplusmesonPtCDFrap = new TH1D("BplusmesonPtCDFrap","BplusmesonPt |y|<1",100,0,50);
  TH1D* BminusmesonPtCDFrap = new TH1D("BminusmesonPtCDFrap","Bminus meson Pt |y|<1",100,0,50);
  TH1D* electronFrombPt = new TH1D("electronFrombPt","electrons from b",100,0,30);
  TH1D* positronFrombPt = new TH1D("positronFrombPt","positrons from b",100,0,30);
  TH1D* epluseminusMinv = new TH1D("epluseminusMinv","e+ e- Inv. Mass",300,0,30);
  TH1D* epluseminusRapidity = new TH1D("epluseminusRapidity","e+ e- y",80,-4,4);
  TH1D* epluseminusMinvMidRap = new TH1D("epluseminusMinvMidRap","e+ e- Inv. Mass |y|<0.5",300,0,30);

  // Begin event loop. Generate event. Skip if error. List first one.
  int nPace = max(1,nEvent/nShow);
  int nErrors(0);
  for (int iEvent = 0; iEvent < nEvent; ++iEvent) {
    if (!pythia.next()) {
      ++nErrors;
      if (nErrors>=nAllowErr) {
    cout << "Reached error limit : " << nErrors << endl;
    cout << "Bailing out! " << endl;
    break;
      }
      continue;
    }
    if (iEvent%nPace == 0) cout << " Now begin event " << iEvent << endl;
    if (iEvent < nList) {pythia.info.list(); pythia.event.list();}
    // Find number of all final charged particles and fill histogram.
    // Find the b (id = 5) and bbar (id = -5), find their daughters,
    // if daughters include electron (id = 11) and positron (id=-11), calculate their
    // invariant mass
    // Status flags:
    //   21 incoming particles of hardest subprocess
    //   23 outgoing particles of hardest subprocess
    //   81-89 primary hadrons produced by hadronization process (B mesons, e.g.)
    //   91-99 particles produced in decay process or by B-E effects (e.g. the electrons)

    int nCharged = 0;
    int indexBQuark(0), indexBbarQuark(0);
    for (int i = 0; i < pythia.event.size(); ++i) {
      if (pythia.event[i].isFinal() && pythia.event[i].isCharged()) {
        ++nCharged;
      }
      Particle& theParticle = pythia.event[i];
   
      if (theParticle.id() == 5 ) {
    indexBQuark = i;
    //cout << "Mother 1, Mother 2 = " << theParticle.mother1() << ", " << theParticle.mother2() << endl;
      }
      if (theParticle.id() == -5) {
    indexBbarQuark = i;
    //cout << "Mother 1, Mother 2 = " << theParticle.mother1() << ", " << theParticle.mother2() << endl;
      }
    } // particle loop

    cout << "Found b quark at index " << indexBQuark << endl;
    cout << "Found bbar quark at index " << indexBbarQuark << endl;
    bquarkPt->Fill(pythia.event[indexBQuark].pT());
    bbarquarkPt->Fill(pythia.event[indexBbarQuark].pT());
    mult.fill( nCharged );
    multHist->Fill(nCharged);
    //cout << "Event " << iEvent << ", Nch= " << nCharged << endl;
   
   
    //Find hadronization products of b and bbar.
    int bQuarkDaughter1 = pythia.event[indexBQuark].daughter1();//first daughter index
    int bQuarkDaughter2 = pythia.event[indexBQuark].daughter2();//last daughter index
    int bbarQuarkDaughter1 = pythia.event[indexBbarQuark].daughter1();
    int bbarQuarkDaughter2 = pythia.event[indexBbarQuark].daughter2();
   
    // Obtain the two hadrons from the fragmentation process
    // Use the PDG id's for this.  All B mesons id's are of the form xx5xx, and
    // all B baryons are of the form 5xxx.
    // So we obtain the id, (make it positive if needed) and then test
    // to see if it is a meson with fmod(currId/100,5)==0.0
    // to see if it is a baryon with currId/1000==5
    int HadronFromBQuark(0), HadronFromBbarQuark(0);
    if (bQuarkDaughter1<bQuarkDaughter2) {
      cout << "Daughters of b Quark" << endl;
      for (int j=bQuarkDaughter1; j<=bQuarkDaughter2; ++j) {
    if (isBHadron(pythia.event[j].id())) {
      cout << "Fragmentation: b -> " << pythia.event[j].name() << endl;
      cout << "                 id " << pythia.event[j].id() << " at index " << j << endl;
      HadronFromBQuark = j;
    }
      }
    }
    if (bbarQuarkDaughter1<bbarQuarkDaughter2) {
      cout << "Daughters of bbar Quark" << endl;
      for (int k=bbarQuarkDaughter1; k<=bbarQuarkDaughter2; ++k) {
    if (isBHadron(pythia.event[k].id())) {
      cout << "Fragmentation : bbar -> " << pythia.event[k].name()  << endl;
      cout << "                     id " << pythia.event[k].id() << " at index " << k << endl;
      HadronFromBbarQuark = k;
    }
      }
    }
    // Search the daughters of the hadrons until electrons and positrons are found
    // if there are any from a semileptonic decay of a beauty hadron
    // Start with the b quark, the b-bar quark loop comes after this
    int Daughter(HadronFromBQuark), electronIndex(0), positronIndex(0);
    while (Daughter!=0) {
      cout << "Checking " << pythia.event[Daughter].name() << " for e+/e- daughters" << endl;
      if (pythia.event[Daughter].id()==-511) {
    // This is a Bbar0, enter its pT
    cout << "Filling Bbar0 pT" << endl;
    B0barmesonPt->Fill(pythia.event[Daughter].pT());
      }
      if (pythia.event[Daughter].id()==511) {
    // This is a B0, enter its pT
    cout << "Filling Bbar0 pT" << endl;
    B0mesonPt->Fill(pythia.event[Daughter].pT());
      }
      Vec4 daughterVec4 = pythia.event[Daughter].p();
      double daughterRap = myRapidity(daughterVec4);

       if (pythia.event[Daughter].id()==-521) {
    // This is a Bminus, enter its pT
    cout << "Filling Bminus pT" << endl;
    BminusmesonPt->Fill(pythia.event[Daughter].pT());
    if (fabs(daughterRap)<1.0) {
      BminusmesonPtCDFrap->Fill(pythia.event[Daughter].pT());
    }
      }
      if (pythia.event[Daughter].id()==521) {
    // This is a Bplus, enter its pT
    cout << "Filling Bplus pT" << endl;
    BplusmesonPt->Fill(pythia.event[Daughter].pT());
    if (fabs(daughterRap)<1.0) {
      BplusmesonPtCDFrap->Fill(pythia.event[Daughter].pT());
    }
      }
     int nextDaughter1 = pythia.event[Daughter].daughter1();
      int nextDaughter2 = pythia.event[Daughter].daughter2();
      // search for electron or positron
      for (int iDaughter = nextDaughter1; iDaughter<=nextDaughter2; ++iDaughter) {
    if (pythia.event[iDaughter].id()==11) {
      cout << "Found electron" << endl;
      cout << pythia.event[iDaughter].name() << endl;
      electronIndex=iDaughter;
      electronFrombPt->Fill(pythia.event[electronIndex].pT());
      break;
    }
    if (pythia.event[iDaughter].id()==-11) {
      cout << "Found positron" << endl;
      cout << pythia.event[iDaughter].name() << endl;
      positronIndex=iDaughter;
      positronFrombPt->Fill(pythia.event[positronIndex].pT());
      break;
    }
      }// loop over daughters to check for e+e-
     
      // If we get here, that means there were no electrons nor positrons.
      // Set the Daughter index to zero now.
      Daughter = 0;
      // If any of the daughters is still a beauty-hadron, we can try again
      // and reset the Daughter index, but only if one of the daughters contains a
      // b quark.
      for (int jDaughter = nextDaughter1; jDaughter<=nextDaughter2; ++jDaughter) {
    if (isBHadron(pythia.event[jDaughter].id())) {
      //One of the daughters is a beauty hadron.
      Daughter = jDaughter;
    }
      }// loop over daughters to check for another b hadron
    }// end of search for electrons in all the daughters of the b quark
   
    // Now search among the daughters of the bbar quark
    Daughter=HadronFromBbarQuark;
    while (Daughter!=0) {
      cout << "Checking " << pythia.event[Daughter].name() << " for e+/e- daughters" << endl;
      if (pythia.event[Daughter].id()==-511) {
    // This is a Bbar0, enter its pT
    cout << "Filling Bbar0 pT" << endl;
    B0barmesonPt->Fill(pythia.event[Daughter].pT());
      }
      if (pythia.event[Daughter].id()==511) {
    // This is a B0, enter its pT
    cout << "Filling B0 pT" << endl;
    B0mesonPt->Fill(pythia.event[Daughter].pT());
      }
      Vec4 daughterVec4 = pythia.event[Daughter].p();
      double daughterRap = myRapidity(daughterVec4);

       if (pythia.event[Daughter].id()==-521) {
    // This is a Bminus, enter its pT
    cout << "Filling Bminus pT" << endl;
    BminusmesonPt->Fill(pythia.event[Daughter].pT());
    if (fabs(daughterRap)<1.0) {
      BminusmesonPtCDFrap->Fill(pythia.event[Daughter].pT());
    }
      }
      if (pythia.event[Daughter].id()==521) {
    // This is a Bplus, enter its pT
    cout << "Filling Bplus pT" << endl;
    BplusmesonPt->Fill(pythia.event[Daughter].pT());
    if (fabs(daughterRap)<1.0) {
      BplusmesonPtCDFrap->Fill(pythia.event[Daughter].pT());
    }
      }

      int nextDaughter1 = pythia.event[Daughter].daughter1();
      int nextDaughter2 = pythia.event[Daughter].daughter2();
      // search for electron or positron
      for (int iDaughter = nextDaughter1; iDaughter<=nextDaughter2; ++iDaughter) {
    //cout << "daughter is a " << pythia.event[iDaughter].name() << endl;
    if (pythia.event[iDaughter].id()==11) {
      cout << "Found electron" << endl;
      cout << pythia.event[iDaughter].name() << endl;
      electronIndex=iDaughter;
      electronFrombPt->Fill(pythia.event[electronIndex].pT());
      break;
    }
    if (pythia.event[iDaughter].id()==-11) {
      cout << "Found positron" << endl;
      cout << pythia.event[iDaughter].name() << endl;
      positronIndex=iDaughter;
      positronFrombPt->Fill(pythia.event[positronIndex].pT());
      break;
    }
      }// loop over daughters to check for e+e-
     
      // If we get here, that means there were no electrons nor positrons.
      // Set the Daughter index to zero now.
      Daughter = 0;
      // If any of the daughters is still a beauty-hadron, we can try again
      // and reset the Daughter index, but only if one of the daughters contains a
      // b quark.
      for (int jDaughter = nextDaughter1; jDaughter<=nextDaughter2; ++jDaughter) {
    if (isBHadron(pythia.event[jDaughter].id())) {
      //One of the daughters is a beauty hadron.
      Daughter = jDaughter;
    }
      }// loop over daughters to check for another b hadron
    }//end of search for electron among daughters of bbar quark
   
    if (electronIndex!=0 && positronIndex!=0) {
      cout << "Found an e+e- pair from bbar" << endl;
      cout << "Ele 4-mom = " << pythia.event[electronIndex].p() << endl;
      cout << "Pos 4-mom = " << pythia.event[positronIndex].p() << endl;
      Vec4 epluseminus(pythia.event[electronIndex].p()+pythia.event[positronIndex].p());
      epluseminusMinv->Fill(epluseminus.mCalc());
      double epluseminusRap = 0.5*log((epluseminus.e()+epluseminus.pz())/(epluseminus.e()-epluseminus.pz()));
      epluseminusRapidity->Fill(epluseminusRap);
      if (fabs(epluseminusRap)<0.5) epluseminusMinvMidRap->Fill(epluseminus.mCalc());
    }
    else {
      cout << "No e+e- pair in event" << endl;
    }
   
  // End of event loop. Statistics. Histogram. Done.
  }// event loop
  if (showStat) pythia.statistics();
  //cout << mult << endl;

  //Write Output ROOT hisotgram into ROOT file
  TFile* outFile = new TFile("pythiaOutputHistosTest.root","RECREATE");
  multHist->Write();
  bquarkPt->Write();
  bbarquarkPt->Write();
  B0mesonPt->Write();
  B0barmesonPt->Write();
  BplusmesonPt->Write();
  BminusmesonPt->Write();
  BplusmesonPtCDFrap->Write();
  BminusmesonPtCDFrap->Write();
  electronFrombPt->Write();
  positronFrombPt->Write();
  epluseminusMinv->Write();
  epluseminusRapidity->Write();
  epluseminusMinvMidRap->Write();
  outFile->Close();

  return 0;
}