// main91.cc is a part of the PYTHIA event generator.
// Copyright (C) 2017 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"

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

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

using namespace Pythia8;

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

  TTree *tree = new TTree("tree","myTree");
  int isBottom20443 = 0;
  int isBottom445 = 0;
  int isBottom10441 = 0;
  int isTop443 = 0;
  int id1, id2;
  gROOT->ProcessLine("#include <vector>");

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

  // Create Pythia instance and set it up to generate hard QCD processes
  // above pTHat = 20 GeV for pp collisions at 14 TeV.
  Pythia pythia;
  pythia.readString("Charmonium:gg2ccbar(3PJ)[3PJ(1)]g = on,on,on");
  pythia.readString("PhaseSpace:pTHatMin = 0");
  pythia.readString("Beams:eCM = 5020.");
  //id:oneChannel = onMode bRatio meMode product1 product2
  pythia.readString("20443:oneChannel= 1 1 0 443 22");
  pythia.readString("445:oneChannel= 1 1 0 443 22");
  pythia.readString("10441:oneChannel= 1 1 0 443 22");
  pythia.readString("443:oneChannel= 1 1 0 13 -13");
  pythia.readString("22:oneChannel= 1 1 0 11 -11");
  pythia.init();

  // Create file on which histogram(s) can be saved.
  TFile* outFile = new TFile("hist1.root", "RECREATE");
  
  
  //float muonPt = 0;
  int muonN = 0;
  Vec4 pvec;
  float m20443, m445, m10441, m443, m22, m443_22, m13_1, m13_2, m13_tot, rap13_1, rap13_2;
  std::vector<float> muonCharge;
  std::vector<float> muonM;
  std::vector<float> muonRap;
  std::vector<float> muonPt;
  std::vector<float> dimuonPt;
  std::vector<float> dimuonRap;
  std::vector<float> dimuonMass;
  std::vector<float> electronRap;
  std::vector<float> electronPt;
  std::vector<float> p3;
  std::vector<float> jpsiPt;
  std::vector<int> isOppSignMuon;
  std::vector<float> gammaEta;
  std::vector<float> gammaPt;
  std::vector<float> electronCharge;

  tree->Branch("muonM", &muonM);
  tree->Branch("muonN", &muonN);
  tree->Branch("muonCharge", &muonCharge);
  tree->Branch("m20443", &m20443);
  tree->Branch("m445", &m445);
  tree->Branch("m10441", &m10441);
  tree->Branch("m443", &m443);
  tree->Branch("m22", &m22);
  tree->Branch("m443_22", &m443_22);
  tree->Branch("m13_1", &m13_1);
  tree->Branch("m13_2", &m13_2);
  tree->Branch("m13_tot", &m13_tot);
  tree->Branch("muonRap", &muonRap);
  tree->Branch("muonPt", &muonPt);
  tree->Branch("dimuonRap", &dimuonRap);
  tree->Branch("dimuonPt", &dimuonPt);
  tree->Branch("dimuonMass", &dimuonMass);
  tree->Branch("electronRap", &electronRap);
  tree->Branch("electronPt", &electronPt);
  tree->Branch("electronCharge", &electronCharge);
  tree->Branch("p3", &p3);
  tree->Branch("jpsiPt", &jpsiPt);
  tree->Branch("isOppSignMuon",&isOppSignMuon);
  tree->Branch("gammaEta", &gammaEta);
  tree->Branch("gammaPt", &gammaPt);

  // Begin event loop. Generate event; skip if generation aborted.
  for (int iEvent = 0; iEvent < 10000; ++iEvent) {
    if(iEvent%100==0){
      cout<<"--Processing event: "<<iEvent<<endl;
    }
    if (!pythia.next()) continue;

    muonN = 0;
    m20443 = 0;
    m445 = 0;
    m10441 = 0;
    m443 = 0;
    m22 = 0;
    m443_22 = 0;
    m13_1 = 0;
    m13_2 = 0;
    m13_tot = 0;
    
    muonM.clear();
    muonCharge.clear();
    muonPt.clear();
    muonRap.clear();
    dimuonPt.clear();
    dimuonRap.clear();
    dimuonMass.clear();
    electronPt.clear();
    electronRap.clear();
    p3.clear();
    isOppSignMuon.clear();
    electronCharge.clear();
    gammaEta.clear();
    gammaPt.clear();

    pvec.reset();

    for (int i = 0; i < pythia.event.size(); ++i){
      if(pythia.event[i].isFinal()&&(pythia.event[i].id()==11 || pythia.event[i].id()==-11)){
        electronPt.push_back(pythia.event[i].pT());
        electronRap.push_back(pythia.event[i].eta());
      }
      if(pythia.event[i].isFinal()&&(pythia.event[i].id()==13 || pythia.event[i].id()==-13)){
        muonN++;
        muonM.push_back(pythia.event[i].m());
        muonCharge.push_back(pythia.event[i].charge());
        muonPt.push_back(pythia.event[i].pT());
        muonRap.push_back(pythia.event[i].eta());
        pvec += pythia.event[i].p();
        //dimuonRap.push_back(pythia.event[i].p()+);
        //dimuonPt.push_back();
      }
      if(pythia.event[i].id()==445 && !isBottom445){
        m445 = pythia.event[pythia.event[i].iBotCopy()].m();
        isBottom445 = 1;
      }
      if(pythia.event[i].id()==20443 && !isBottom20443){
        m20443 = pythia.event[pythia.event[i].iBotCopy()].m();
        isBottom20443 = 1;
      }
      if(pythia.event[i].id()==10441 && !isBottom10441){
        m10441 = pythia.event[pythia.event[i].iBotCopy()].m();
        isBottom10441 = 1;
      }
      id1 = pythia.event[pythia.event[i].iBotCopy()].daughter1();
      id2 = pythia.event[pythia.event[i].iBotCopy()].daughter2();
      m443_22 = (pythia.event[id1].p()+pythia.event[id2].p()).mCalc();
      if(pythia.event[id1].id() == 443){
        m443 = pythia.event[id1].m();
        m22 = pythia.event[id2].m();
        m13_1 = pythia.event[pythia.event[pythia.event[id1].iBotCopy()].daughter1()].m();
        m13_2 = pythia.event[pythia.event[pythia.event[id1].iBotCopy()].daughter2()].m();
        m13_tot = (pythia.event[pythia.event[pythia.event[id1].iBotCopy()].daughter1()].p()+pythia.event[pythia.event[pythia.event[id1].iBotCopy()].daughter2()].p()).mCalc();
        //jpsiPt.push_back(pythia.event[id1].pT());
      }else{
        m443 = pythia.event[id2].m();
        m22 = pythia.event[id1].m();
        m13_1 = pythia.event[pythia.event[pythia.event[id2].iBotCopy()].daughter1()].m();
        m13_2 = pythia.event[pythia.event[pythia.event[id2].iBotCopy()].daughter2()].m();
        m13_tot = (pythia.event[pythia.event[pythia.event[id2].iBotCopy()].daughter1()].p()+pythia.event[pythia.event[pythia.event[id2].iBotCopy()].daughter2()].p()).mCalc();
        jpsiPt.push_back(pythia.event[id2].pT());
      }
    }
    if(muonCharge[0]!=muonCharge[1]){
      isOppSignMuon.push_back(1);
    }else{
      isOppSignMuon.push_back(0);
    }
    dimuonPt.push_back(pvec.pT());
    dimuonRap.push_back(pvec.rap());
    p3.push_back(pvec.pAbs());
    dimuonMass.push_back(pvec.mCalc());
    tree->Fill();
  }

  // Statistics on event generation.
  pythia.stat();


  // Save histogram on file and close file.
  tree->Write();

  delete outFile;

  // Done.
  return 0;
  //cout<<"HELLO"<<endl;
}