// main91.cc is a part of the PYTHIA event generator.
// Copyright (C) 2016 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.
// It studies the charged multiplicity distribution at the LHC.
// Modified by Rene Brun, Axel Naumann and Bernhard Meirose
// to use ROOT for histogramming.

// Stdlib header file for input and output.
#include <iostream>

// Header file to access Pythia 8 program elements.
#include "Pythia8/Pythia.h"
// ROOT, for histogramming.
#include "TH1.h"
#include "TH2.h"

#include<TCanvas.h>
#include<TFrame.h>

// ROOT, for interactive graphics.
#include "TVirtualPad.h"
#include "TApplication.h"

// ROOT, for saving file.
#include "TFile.h"

using namespace Pythia8;


//==========================================================================
// A derived class to do J/psi decays.

class JpsiDecay : public DecayHandler {

public:

  // Constructor.
 JpsiDecay(ParticleData* pdtPtrIn, Rndm* rndmPtrIn) {times = 0;
 pdtPtr = pdtPtrIn; rndmPtr = rndmPtrIn;}

  // Routine for doing the decay.
bool decay(vector<int>& idProd, vector<double>& mProd,
vector<Vec4>& pProd, int iDec, const Event& event);

private:

  // Count number of times JpsiDecay is called.
int times;

  //Pointer to the particle data table.
ParticleData* pdtPtr;

  // Pointer to the random number generator.
 Rndm* rndmPtr;

};

//--------------------------------------------------------------------------

// The actual J/psi decay routine.
// Not intended for realism, just to illustrate the principles.

bool JpsiDecay::decay(vector<int>& idProd, vector<double>& mProd,
vector<Vec4>& pProd, int iDec, const Event& event) {

// Always do decay J/psi -> mu+ mu-; store the muons.
idProd.push_back(22);
idProd.push_back(22);

  // Muon mass(es), here from Pythia tables, also stored.
double mMuon = pdtPtr->m0(22);
mProd.push_back(mMuon);
mProd.push_back(mMuon);

  // Calculate muon energy and momentum in J/psi rest frame.
double eMuon = 0.5 * mProd[0];
double pAbsMuon = sqrt(eMuon * eMuon - mMuon * mMuon);

  // Assume decay angles isotropic in rest frame.
 double cosTheta = 2. * rndmPtr->flat() - 1.;
double sinTheta = sqrt(max(0., 1. - cosTheta * cosTheta));
double phi = 2. * M_PI * rndmPtr->flat();
double pxMuon = pAbsMuon * sinTheta * cos(phi);
double pyMuon = pAbsMuon * sinTheta * sin(phi);
double pzMuon = pAbsMuon * cosTheta;

  // Define mu+ and mu- four-vectors in the J/psi rest frame.
Vec4 pMuPlus(   pxMuon,  pyMuon,  pzMuon, eMuon);
Vec4 pMuMinus( -pxMuon, -pyMuon, -pzMuon, eMuon);

  // Boost them by velocity vector of the J/psi mother and store.
 pMuPlus.bst(pProd[0]);
 pMuMinus.bst(pProd[0]);
 pProd.push_back(pMuPlus);
 pProd.push_back(pMuMinus);

  // Print message the first few times, to show that it works.
 if (times++ < 10) {
  int iMother = event[iDec].mother1();
  int idMother = event[iMother].id();
 cout << "\n J/psi decay performed, J/psi in line " << iDec
   << ", mother id = " << idMother << "\n";
 }

  // Done
 return true;

}

//==========================================================================



int main(int argc, char* argv[]) {

  // Create the ROOT application environment.
   TApplication theApp("hist24", &argc, argv);

  // Create Pythia instance and set it up to generate hard QCD processes
  // above pTHat = 20 GeV for pp collisions at 13 TeV.
  Pythia pythia;
  
  pythia.readString("HardQCD:all = on");
  pythia.readString("PhaseSpace:pTHatMin = 2.0");  
  pythia.readString("PhaseSpace:pTHatMinDiverge = 0.5"); 
  pythia.readString("PartonLevel:MPI = on");
  pythia.readString("Charmonium:all = on");
  pythia.readString("Tune:pp = 5");
  pythia.readString("Beams:eCM = 13000.");
  pythia.readString("Beams:idA = 2212");
  pythia.readString("Beams:idB = 2212");
  pythia.settings.listAll();



  //  UserHooks* oniumUserHook = new SuppressSmallPT( 1., 3, false);
  // pythia.setUserHooksPtr( oniumUserHook);

  // A class to do J/psi decays externally.
  // DecayHandler* handleDecays = new JpsiDecay(&pythia.particleData,
  // &pythia.rndm);

  // The list of particles the class can handle.
  // vector<int> handledParticles;
  // handledParticles.push_back(443);

  // Hand pointer and list to Pythia.
  //  pythia.setDecayPtr( handleDecays, handledParticles);

  // Switch off automatic event listing in favour of manual.
  pythia.readString("Next:numberShowInfo = 0");
  pythia.readString("Next:numberShowProcess = 0");
  pythia.readString("Next:numberShowEvent = 0");

  // Initialization.
  pythia.init();


     
  // Create file on which histogram(s) can be saved.
  TFile* outFile = new TFile("hist24.root", "RECREATE");
  //TH2F *corr = new TH2F("corr","Correlation", 10,0,500,20, 0.,20.);
  //  TH1F *corr = new TH1F("corr","Correlation", 500.0,0,500.0);

  // Book histogram.
  // TH1F *mult = new TH1F("mult","charged multiplicity", 500.0,0,500.0);
  TH1F *jpsi = new TH1F("jpsi","jpsi multiplicity",200, 0., 5.);
  TH1F *pTJPsi=new TH1F("pTJPsi","pT of J/Psi", 50, 0., 50.);

  //===========Mult Bin 0-1=================
  //  TH1F *mult_0_1 = new TH1F("mult_0_1","charged multiplicity", 500,0.0,500.0);
  TH1F *jpsi_0_1 = new TH1F("jpsi_0_1","jpsi multiplicity",200, 0., 5.);
  TH1F *pTJPsi_0_1=new TH1F("pTJPsi_0_1","pT of J/Psi", 50, 0., 50.);
  //===========Mult Bin 1-2=================
  // TH1F *mult_1_2 = new TH1F("mult_1_2","charged multiplicity",500,0.0,500.0);
  TH1F *jpsi_1_2 = new TH1F("jpsi_1_2","jpsi multiplicity",200, 0., 5.);
  TH1F *pTJPsi_1_2=new TH1F("pTJPsi_1_2","pT of J/Psi", 50, 0., 50.);
  //===========Mult Bin 2-3=================
  // TH1F *mult_2_3 = new TH1F("mult_2_3","charged multiplicity",500,0.0,500.0);
  TH1F *jpsi_2_3 = new TH1F("jpsi_2_3","jpsi multiplicity",200, 0., 5.);
  TH1F *pTJPsi_2_3=new TH1F("pTJPsi_2_3","pT of J/Psi", 50, 0., 50.);
  //===========Mult Bin 3-4=================
  //  TH1F *mult_3_4 = new TH1F("mult_3_4","charged multiplicity",500,0.0,500.0);
  TH1F *jpsi_3_4 = new TH1F("jpsi_3_4","jpsi multiplicity",200, 0., 5.);
  TH1F *pTJPsi_3_4=new TH1F("pTJPsi_3_4","pT of J/Psi", 50, 0., 50.);
  //===========Mult Bin 4-5=================
  // TH1F *mult_4_5 = new TH1F("mult_4_5","charged multiplicity",500,0.0,500.0);
  TH1F *jpsi_4_5 = new TH1F("jpsi_4_5","jpsi multiplicity",200, 0., 5.);
  TH1F *pTJPsi_4_5=new TH1F("pTJPsi_4_5","pT of J/Psi", 50, 0., 50.);
  //===========Mult Bin 5-6=================
  // TH1F *mult_5_6 = new TH1F("mult_5_6","charged multiplicity",500,0.0,500.0);
  TH1F *jpsi_5_6 = new TH1F("jpsi_5_6","jpsi multiplicity",200, 0., 5.);
  TH1F *pTJPsi_5_6=new TH1F("pTJPsi_5_6","pT of J/Psi", 50, 0., 50.);
  //===========Mult Bin 6-7=================
  // TH1F *mult_6_7 = new TH1F("mult_6_7","charged multiplicity",500,0.0,500.0);
  TH1F *jpsi_6_7 = new TH1F("jpsi_6_7","jpsi multiplicity",200, 0., 5.);
  TH1F *pTJPsi_6_7=new TH1F("pTJPsi_6_7","pT of J/Psi", 50, 0., 50.);
  //===========Mult Bin 7-8=================
  // TH1F *mult_7_8 = new TH1F("mult_7_8","charged multiplicity",500,0.0,500.0);
  TH1F *jpsi_7_8 = new TH1F("jpsi_7_8","jpsi multiplicity",200, 0., 5.);
  TH1F *pTJPsi_7_8=new TH1F("pTJPsi_7_8","pT of J/Psi", 50, 0., 50.);
  //===========Mult Bin 8-9=================
  // TH1F *mult_8_9 = new TH1F("mult_8_9","charged multiplicity",500,0.0,500.0);
  TH1F *jpsi_8_9 = new TH1F("jpsi_8_9","jpsi multiplicity",200, 0., 5.);
  TH1F *pTJPsi_8_9=new TH1F("pTJPsi_8_9","pT of J/Psi", 50, 0., 50.);
  //===========Mult Bin 9-10=================
  // TH1F *mult_9_10 = new TH1F("mult_9_10","charged multiplicity",500,0.0,500.0);
  TH1F *jpsi_9_10 = new TH1F("jpsi_9_10","jpsi multiplicity",200, 0., 5.);
  TH1F *pTJPsi_9_10=new TH1F("pTJPsi_9_10","pT of J/Psi", 50, 0., 50.);

  // Begin event loop. Generate event; skip if generation aborted.
  Double_t pt_jpsi; 
  for (int iEvent = 0; iEvent < 10000 ; ++iEvent) {
    if (!pythia.next()) continue;

    //Number of J/Psi
     int Njpsi=0;
     int nCharged = 0;
     Double_t pt;
     for (int i = 0; i < pythia.event.size(); ++i)
       {
      
    // Find number of all final charged particles.
       Double_t Eta = pythia.event[i].y();
       if(Eta >-1.0 && Eta < 1.0)
	 {
      if (pythia.event[i].isFinal() && pythia.event[i].isCharged())
      ++nCharged;
	 }

       if(Eta >-0.5 && Eta < 0.5)
	 {
       if (pythia.event[i].id() == 443)  
       ++Njpsi;
       if (pythia.event[i].id() == 443) pt_jpsi = pythia.event[i].pT();
       if (pythia.event[i].id() == 443) pTJPsi->Fill( pythia.event[i].pT() );
	 }

       }

     // mult->Fill(nCharged,nCharged>0);
     // jpsi->Fill(Njpsi,Njpsi>0);
    // corr->Fill(nCharged,Njpsi);
    
     
    if(nCharged >=0 && nCharged <=1)
      {
	// mult_0_10->Fill(nCharged,nCharged>0);
	// jpsi_0_1->Fill(Njpsi,Njpsi>0);
      pTJPsi_0_1->Fill(pt_jpsi);
      }

    if(nCharged >=1 && nCharged <=2)
      {
	//	mult_10_20->Fill(nCharged,nCharged>0);
	// jpsi_1_2->Fill(Njpsi,Njpsi>0);
      pTJPsi_1_2->Fill(pt_jpsi);
      }

    if(nCharged >=2 && nCharged <=3)
      {
	// mult_20_40->Fill(nCharged,nCharged>0);
	// jpsi_2_3->Fill(Njpsi,Njpsi>0);
      pTJPsi_2_3->Fill(pt_jpsi);
      }

     if(nCharged >=3 && nCharged <=4)
      {
	// mult_40_60->Fill(nCharged,nCharged>0);
	// jpsi_3_4->Fill(Njpsi,Njpsi>0);
      pTJPsi_3_4->Fill(pt_jpsi);
      }

      if(nCharged >=4 && nCharged <=5)
      {
	// mult_60_120->Fill(nCharged,nCharged>0);
	// jpsi_4_5->Fill(Njpsi,Njpsi>0);
      pTJPsi_4_5->Fill(pt_jpsi);
      }

       if(nCharged >=5 && nCharged <=6)
      {
	// mult_120_200->Fill(nCharged,nCharged>0);
	// jpsi_5_6->Fill(Njpsi,Njpsi>0);
      pTJPsi_5_6->Fill(pt_jpsi);
      }

        if(nCharged >=6 && nCharged <=7)
      {
	// mult_200_500->Fill(nCharged,nCharged>0);
	// jpsi_6_7->Fill(Njpsi,Njpsi>0);
       pTJPsi_6_7->Fill(pt_jpsi);
      }

	 if(nCharged >=7 && nCharged <=8)
      {
	// mult_500_700->Fill(nCharged,nCharged>0);
	// jpsi_7_8->Fill(Njpsi,Njpsi>0);
       pTJPsi_7_8->Fill(pt_jpsi);
      }

       if(nCharged >=8 && nCharged <=9)
      {
	// mult_700_1000->Fill(nCharged,nCharged>0);
	// jpsi_8_9->Fill(Njpsi,Njpsi>0);
       pTJPsi_8_9->Fill(pt_jpsi);
      }

      if(nCharged >=9 && nCharged <=10)
      {
	// mult_1000_1500->Fill(nCharged,nCharged>0);
	// jpsi_9_10->Fill(Njpsi,Njpsi>0);
       pTJPsi_9_10->Fill(pt_jpsi);

      pTJPsi_0_1->Draw();
      pTJPsi_1_2->Draw("SAME");
      pTJPsi_2_3->Draw("SAME");
      pTJPsi_3_4->Draw("SAME");
      pTJPsi_4_5->Draw("SAME");
      pTJPsi_5_6->Draw("SAME");
      pTJPsi_6_7->Draw("SAME");
      pTJPsi_7_8->Draw("SAME");
      pTJPsi_8_9->Draw("SAME");
      pTJPsi_9_10->Draw("SAME");
      
      }
    
  }
  
 
  // Statistics on event generation.
  pythia.stat();

  // Show histogram. Possibility to close it.
  //  mult->Draw();
  // std::cout << "\nDouble click on the histogram window to quit.\n";
  // gPad->WaitPrimitive();

  // Save histogram on file and close file.
  // mult->Write();
  // jpsi->Write();
  // corr->Write();
  pTJPsi->Write();

  // mult_0_1->Write();
  // jpsi_0_1->Write();
  pTJPsi_0_1->Write();

   //  mult_1_2->Write();
   // jpsi_1_2->Write();
   pTJPsi_1_2->Write();
  

  // mult_2_3->Write();
  //  jpsi_2_3->Write();
  pTJPsi_2_3->Write();

  //  mult_3_4->Write();
  // jpsi_3_4->Write();
  pTJPsi_3_4->Write();

  // mult_4_5->Write();
  // jpsi_4_5->Write();
  pTJPsi_4_5->Write();

  // mult_5_6->Write();
  // jpsi_5_6->Write();
  pTJPsi_5_6->Write();

  // mult_6_7->Write();
  // jpsi_6_7->Write();
  pTJPsi_6_7->Write();
   
  // mult_7_8->Write();
  //jpsi_7_8->Write();
  pTJPsi_7_8->Write();
  
  // mult_8_9->Write();
  //jpsi_8_9->Write();
   pTJPsi_8_9->Write();
  
  //mult_9_10->Write();
   //psi_9_10->Write();
  pTJPsi_9_10->Write();
  
  
 
  //delete outFile;

  // Done.
  return 0;

   
}