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

#include "fastjet/PseudoJet.hh"
#include "TTree.h"
#include "TFile.h"
#include "TBranch.h"

const int kMaxEvent = 1;
const int kMaxParticle = 20; // maximum number of jets?
const int kMaxRwgt = 1;

class outputTree {

  TTree* fTree;
  TFile* treeFile;

public:

  // Declaration of branches & leaves in a C-type struct
  // See: https://root.cern.ch/root/htmldoc/guides/users-guide/ROOTUsersGuide.html (12.15 Example 2: A Tree with a C Structure)
  // See: https://root-forum.cern.ch/t/ttree-with-c-struct/13469/2
  typedef struct {
    int           Event_;
    UInt_t        fUniqueID[kMaxEvent];   //[Event_]
    UInt_t        fBits[kMaxEvent];   //[Event_]
    Long64_t      Number[kMaxEvent];   //[Event_]
    int           Nparticles[kMaxEvent];   //[Event_]
    int           ProcessID[kMaxEvent];   //[Event_]
    double        Weight[kMaxEvent];   //[Event_]
    double        ScalePDF[kMaxEvent];   //[Event_]
    double        CouplingQED[kMaxEvent];   //[Event_]
    double        CouplingQCD[kMaxEvent];   //[Event_]
    int           size;
  } EVENT;
 
  typedef struct {
    int           Rwgt_;
    UInt_t        fUniqueID[kMaxRwgt];   //[Rwgt_]
    UInt_t        fBits[kMaxRwgt];   //[Rwgt_]
    Double_t      Weight[kMaxRwgt];   //[Rwgt_]
    int           size;
  } RWGT;

  typedef struct {
    int           Particle_;
    UInt_t        fUniqueID[kMaxParticle];   //[Particle_]
    UInt_t        fBits[kMaxParticle];   //[Particle_]
    int           PID[kMaxParticle];   //[Particle_]
    int           Status[kMaxParticle];   //[Particle_]
    int           Mother1[kMaxParticle];   //[Particle_]
    int           Mother2[kMaxParticle];   //[Particle_]
    int           ColorLine1[kMaxParticle];   //[Particle_]
    int           ColorLine2[kMaxParticle];   //[Particle_]
    double        Px[kMaxParticle];   //[Particle_]
    double        Py[kMaxParticle];   //[Particle_]
    double        Pz[kMaxParticle];   //[Particle_]
    double        E[kMaxParticle];   //[Particle_]
    double        M[kMaxParticle];   //[Particle_]
    double        PT[kMaxParticle];   //[Particle_]
    double        Eta[kMaxParticle];   //[Particle_]
    double        Phi[kMaxParticle];   //[Particle_]
    double        Rapidity[kMaxParticle];   //[Particle_]
    double        LifeTime[kMaxParticle];   //[Particle_]
    double        Spin[kMaxParticle];   //[Particle_]
    int           size;
  } PARTICLE;

  EVENT Event;
  RWGT Rwgt;
  PARTICLE Particle;

  // Member functions

  void initTree( TString&, double );
  //void clearValues();
  void fillTree( vector <fastjet::PseudoJet>, int );
  void writeTree();

};

void outputTree::initTree( TString& inputLHE, double pTjetmin )
{  

  char outTree[500];
  sprintf(outTree, "%s_jetsreco_pTjetmin%iGeV.root", inputLHE.Data(), (int)pTjetmin);
  treeFile = new TFile(outTree, "recreate");
  fTree = new TTree("LHEF", "LHEF (reco jets)");
  fTree->SetDirectory(treeFile);
  fTree->SetAutoSave(100000);



  // Set branch & leaves addresses 

  fTree->Branch("Event_",&Event.Event_,"Event_/I");
  fTree->Branch("Event_fUniqueID",&Event.fUniqueID,"Event_fUniqueID[Event_]/i");
  fTree->Branch("Event_fBits",&Event.fBits,"Event_fBits[Event_]/i");
  fTree->Branch("Event_Number",&Event.Number,"Event_Number[Event_]/L");
  fTree->Branch("Event_Nparticles",&Event.Nparticles,"Event_Nparticles[Event_]/I");
  fTree->Branch("Event_ProcessID",&Event.ProcessID,"Event_ProcessID[Event_]/I");
  fTree->Branch("Event_Weight",&Event.Weight,"Event_Weight[Event_]/D");
  fTree->Branch("Event_ScalePDF",&Event.ScalePDF,"Event_ScalePDF[Event_]/D");
  fTree->Branch("Event_CouplingQED",&Event.CouplingQED,"Event_CouplingQED[Event_]/D");
  fTree->Branch("Event_CouplingQCD",&Event.CouplingQCD,"Event_CouplingQCD[Event_]/D");
  fTree->Branch("Event_size",&Event.size, "Event_size/I");

  fTree->Branch("Rwgt_",&Rwgt.Rwgt_,"Rwgt_/I");
  fTree->Branch("Rwgt_fUniqueID",&Rwgt.fUniqueID,"Rwgt_fUniqueID[Rwgt_]/i");
  fTree->Branch("Rwgt_fBits",&Rwgt.fBits,"Rwgt_fBits[Rwgt_]/i");
  fTree->Branch("Rwgt_Weight",&Rwgt.Weight,"Rwgt_Weight[Rwgt_]/D");
  fTree->Branch("Rwgt_size",&Rwgt.size,"Rwgt_size/I");

  fTree->Branch("Particle_",&Particle.Particle_,"Particle_/I");
  fTree->Branch("Particle_fUniqueID",&Particle.fUniqueID,"Particle_fUniqueID[Particle_]/i");
  fTree->Branch("Particle_fBits",&Particle.fBits,"Particle_fBits[Particle_]/i");
  fTree->Branch("Particle_PID",&Particle.PID,"Particle_PID[Particle_]/I");
  fTree->Branch("Particle_Status",&Particle.Status,"Particle_Status[Particle_]/I");
  fTree->Branch("Particle_Mother1",&Particle.Mother1,"Particle_Mother1[Particle_]/I");
  fTree->Branch("Particle_Mother2",&Particle.Mother2,"Particle_Mother2[Particle_]/I");
  fTree->Branch("Particle_ColorLine1",&Particle.ColorLine1,"Particle_ColorLine1[Particle_]/I");
  fTree->Branch("Particle_ColorLine2",&Particle.ColorLine2,"Particle_ColorLine2[Particle_]/I");
  fTree->Branch("Particle_Px",&Particle.Px,"Particle_Px[Particle_]/D");
  fTree->Branch("Particle_Py",&Particle.Py,"Particle_Py[Particle_]/D");
  fTree->Branch("Particle_Pz",&Particle.Pz,"Particle_Pz[Particle_]/D");
  fTree->Branch("Particle_E",&Particle.E,"Particle_E[Particle_]/D");
  fTree->Branch("Particle_M",&Particle.M,"Particle_M[Particle_]/D");
  fTree->Branch("Particle_PT",&Particle.PT,"Particle_PT[Particle_]/D");
  fTree->Branch("Particle_Eta",&Particle.Eta,"Particle_Eta[Particle_]/D");
  fTree->Branch("Particle_Phi",&Particle.Phi,"Particle_Phi[Particle_]/D");
  fTree->Branch("Particle_Rapidity",&Particle.Rapidity,"Particle_Rapidity[Particle_]/D");
  fTree->Branch("Particle_LifeTime",&Particle.LifeTime,"Particle_LifeTime[Particle_]/D");
  fTree->Branch("Particle_Spin",&Particle.Spin,"Particle_Spin[Particle_]/D");
  fTree->Branch("Particle_size",&Particle.size,"Particle_size/I");

}

void outputTree::fillTree( vector <fastjet::PseudoJet> jets, int evtnum )
{

  // Per event quantities
  // Not used, so set to zero...
  Event.Event_ = 1;
  Event.fUniqueID[0] = 0;
  Event.fBits[0]= 0;
  Event.Number[0] = evtnum;
  Event.Nparticles[0]= jets.size();
  Event.ProcessID[0]= 0;
  Event.Weight[0]= 0;
  Event.ScalePDF[0]= 0;
  Event.CouplingQED[0] = 0;
  Event.CouplingQCD[0] = 0;
  Event.size = 1;

  // Per weight quantities
  // Not used, so set to zero...
  Rwgt.Rwgt_ = 1;
  Rwgt.fUniqueID[0] = 0;
  Rwgt.fBits[0] = 0;
  Rwgt.Weight[0] = 0;
  Rwgt.size = 1;

  // Per jet quantities

  Particle.Particle_ = jets.size();

  cout << endl << " *************************************** JET SIZE = " << jets.size() <<  " **************************************" << endl;

  for ( UInt_t i = 0; i < jets.size(); i++) {
   
    Particle.fUniqueID[i] = 0;
    Particle.fBits[i] = 0;
    Particle.PID[i] = 0;
    Particle.Status[i] = 1; // Final-state particles
    Particle.Mother1[i] = 0;
    Particle.Mother2[i] = 0;
    Particle.ColorLine1[i] = 0;
    Particle.ColorLine2[i] = 0;

    Particle.Px[i] = jets[i].px();
    Particle.Py[i] = jets[i].py();
    Particle.Pz[i] = jets[i].px();
    Particle.E[i] = jets[i].E();
    Particle.M[i] = jets[i].m();
    Particle.PT[i] = jets[i].perp();
    Particle.Eta[i] = jets[i].eta();
    Particle.Phi[i] = jets[i].phi_std();
    Particle.Rapidity[i] = jets[i].rap();

    Particle.LifeTime[i] = 0;
    Particle.Spin[i] = 0;
  }
  Particle.size = jets.size();

  fTree->Fill();

}

void outputTree::writeTree(){

  cout << endl << endl;
//  fTree->Print();
  fTree->Write();
  treeFile->Close();

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

}