/*
  Author: David d'Enterria CERN (dde@cern.ch). 2023

  Purpose: Open an LHE file, shower and hadronize its contents with PYTHIA8. Do jet reco. Save an output ROOT file with the TTree jet info.
  //gSystem->Load("/home/enterria/progs/fastjet/install/lib/libsiscone.so");
  //gSystem->Load("/home/enterria/progs/fastjet/install/lib/libsiscone_spherical.so");
  //gSystem->Load("/home/enterria/progs/fastjet/install/lib/libfastjetplugins.so");

  gInterpreter->AddIncludePath("-I$PYTHIA8/include/");
  gSystem->Load("$PYTHIA8/lib/libpythia8");
  gSystem->Load("libEG");
  gSystem->Load("libEGPythia8");

  gInterpreter->AddIncludePath("-I$HOME/progs/fastjet/install/include");
  gSystem->Load("$HOME/progs/fastjet/install/lib/libfastjet.so");




  .L pythia8_shower_hadronize_pp_Njets_from_lhe.C++

  pythia8_shower_hadronize_pp_Njets_from_lhe("events", 35., 3);

.> logfile.txt
 pythia8_shower_hadronize_pp_Njets_from_lhe("pp6j_test", 35., 3);
.>

*/

#include <vector>
#include <iostream>
#include <iomanip>
#include <string>

//#include "TFile.h"
//#include "TTree.h"
//#include "TBranch.h"
#include "TH1D.h"
#include "TRandom3.h"
#include "TClonesArray.h"

#include "fastjet/PseudoJet.hh"
#include "fastjet/ClusterSequence.hh"
//#include "fastjet/JadePlugin.hh"

#include "TPythia8.h"
#include "TParticle.h"
#include "Pythia8/Pythia.h"
#include "Pythia8/Info.h"
#include "Pythia8/Analysis.h"
#include "Pythia8/Settings.h"

#include "pythia8_shower_hadronize_pp_Njets_from_lhe.h"

void pythia8_shower_hadronize_pp_Njets_from_lhe( TString inputLHEfile = "pp6j_LO_alpgen", const double pTjetmin = 35., const int nev = 1e6 )
{

  char text[1000];

  //_____________________________________________________________________________
  // Setup    

  // Define output tree with Jet variables

  outputTree outTree;
  outTree.initTree(inputLHEfile,pTjetmin);

  // create the Pythia object
  TPythia8 *pythia8 = new TPythia8();
  //pythia8->ReadString("Random:setSeed = on");
  //char ranstr[500];
  //sprintf(ranstr, "Random:seed = %d", 10070002+num);
  //pythia8->ReadString(ranstr);

  // PYTHIA-8 settings
  //pythia8->ReadString("PartonLevel:ISR = off"); // NO g/photon ISR
  //pythia8->ReadString("PartonLevel:FSR = off"); // NO g/photon FSR
  //pythia8->ReadString("PartonLevel:MPI = off");
  //pythia8->ReadString("HadronLevel:Hadronize = off"); 
  //pythia8->ReadString("HadronLevel:Decay = on"); // SWITCH-ON HADRON DECAYS
  // Attempt to avoid checking errors in the event
  //pythia8->ReadString(" Check:event = off");

  // Select Les Houches Event File (LHEF) run.

  pythia8->ReadString("Next:numberShowLHA = 0");
  pythia8->ReadString("Next:numberShowInfo = 0");
  pythia8->ReadString("Next:numberShowProcess = 0");
  pythia8->ReadString("Next:numberShowEvent = 0");
  pythia8->ReadString("Beams:frameType = 4");

  // Read the input LHE file
  char lheset[300];
  sprintf(lheset, "Beams:LHEF = %s.lhe",inputLHEfile.Data());
  pythia8->ReadString(lheset);
//  pythia8->ReadString("Beams:LHEF=pp_3j_NLO_H_T_2_35GeV.lhe");

  // Array of particles per event
  TClonesArray* particles = new TClonesArray("TParticle", 10000);
  
  // Fastjet setup
  // Select algorithm and parameters
  std::vector <fastjet::PseudoJet> fjInputs;
  double Rparam = 0.4; // distance parameter = 0.4
  fastjet::JetDefinition *jetDef = new fastjet::JetDefinition(fastjet::antikt_algorithm,Rparam);

  // Event loop

  for (int iev = 0; iev < nev; iev++) {

    if (iev==0) cout << "Generating event #: " << flush;
    if (iev%100 == 0) cout << iev << " " << flush;

    // Get the full event object to initialize it. It's very verbose but the only way the LHEF file seems to be properly initialized...
    Pythia8::Pythia *pythia8full = pythia8->Pythia8();
    if (iev==0) pythia8full->init();

    //Clear the tree variables
    //outTree.clearValues();


      pythia8->GenerateEvent();
      if (iev < 1) pythia8full->LHAeventList();
      pythia8->ImportParticles(particles,"All");
      Int_t np = particles->GetEntriesFast();

     // Reset Fastjet input
    fjInputs.resize(0);

     // Loop over event record to decide what to pass to FastJet
    for (int i = 0; i < np; ++i) {
      TParticle* part = (TParticle*) particles->At(i);
      int ist = part->GetStatusCode();
      // Positive codes are final particles.
      if (ist <= 0) continue;
      int pdg = TMath::Abs(part->GetPdgCode());

      // No neutrinos
      if (pdg == 12 || pdg == 14 || pdg == 16) continue;

      /*
      // Fill histos
      float eta = part->Eta();
      hdNdeta->Fill(eta);
      float pT  = part->Pt();
      hdNdpT->Fill(pT);
      //if (pT > 0.) hdNdpT->Fill(pT, 1./(2.*pT));
      */

      // build jets from particles
      float px  = part->Px();
      float py  = part->Py();
      float pz  = part->Pz();
      float energy  = part->Energy();

      // Only |eta| < 3.6
      //if (TMath::Abs(eta) > 3.6) continue;

      // Store as input to Fastjet
      fjInputs.push_back( fastjet::PseudoJet( px, py, pz, energy ) );
//      cout << "vals: " <<px << " "<< py << " " << pz << " " << energy << endl;

    } // end PARTICLES LOOP

   if (fjInputs.size() == 0) {
      cout << "Error: event with no final state particles clustered into jets" << endl;
      continue;
    }

    // Run Fastjet algorithm
    vector <fastjet::PseudoJet> inclusiveJets, sortedJets;
    fastjet::ClusterSequence clustSeq(fjInputs, *jetDef);

    // For the first event, print the FastJet details
    if (iev<1) {
      cout << "Ran " << jetDef->description() << endl;
      cout << "Strategy adopted by FastJet was "
           << clustSeq.strategy_string() << endl << endl;
    }

    // Extract inclusive jets sorted by pT 

    inclusiveJets = clustSeq.inclusive_jets(pTjetmin);
    sortedJets    = sorted_by_pt(inclusiveJets);

    cout << "Number of reconstructed jets (with pT(j)  > " << int(pTjetmin) << " GeV): " << sortedJets.size() << endl;
    
    // Fill output tree
    outTree.fillTree(sortedJets, iev);

    /*
    // Fill histos
    for (unsigned int i = 0; i < sortedJets.size(); i++) {
      // Only count jets that have |eta| < 2.0
      //if (fabs(sortedJets[i].rap()) > 2.0) continue;
      double etajet = sortedJets[i].rap();
      hdNjetdeta->Fill(etajet);
      double pTjet = sortedJets[i].perp();
      hdNjetdpT->Fill(pTjet);
      //if (pTjet > 0.) hdNjetdpT->Fill(pTjet, 1./(2.*pTjet))
    }
    */
  } // end EVENT LOOP
  
  pythia8->PrintStatistics();

  outTree.writeTree();
  
  /*
  // Output histos file
  char outfile[500];
  sprintf(outfile,"histos_pythia8_fastjet.root");
  TFile f(outfile,"RECREATE");
  hdNdpT->Write();
  hdNjetdpT->Write();
  hdNdeta->Write();
  hdNjetdeta->Write();
  f.Close();

  cout << endl << " <I> File " << outfile << " created ..." << endl;
  cout << " <I> Check contents via:   TFile *ff = TFile::Open(\"" << outfile << "\");" << endl;
  cout << "                           TBrowser b" << endl;
  */

}