gSystem->Load("libfastjet");

%%cpp -d
#include "fastjet/ClusterSequence.hh"
void fillInputParticles(vector<fastjet::PseudoJet>& input_particles, const char* inputFileName) 
{
  auto ntupleFormat = "px:py:pz:E";
  TNtuple input_particles_ntuple("InputParticles","Input Particles",ntupleFormat);
  auto n_particles = input_particles_ntuple.ReadFile(inputFileName,ntupleFormat);
  input_particles.reserve(n_particles);

  for (auto i : ROOT::TSeqI(n_particles)) {
     input_particles_ntuple.GetEntry(i);
     auto v = input_particles_ntuple.GetArgs();
     fastjet::PseudoJet particle(v[0],v[1],v[2],v[3]);
     particle.set_user_index(i);
     input_particles.emplace_back(particle); 
  }
}

%%cpp -d
THStack* getJetsComponents(const vector<fastjet::PseudoJet>& input_particles,
                          fastjet::JetAlgorithm jetAlgorithm = fastjet::antikt_algorithm, 
                          double R = .6,
                          double ptMin = 14.)
{
   // create a jet definition: 
   // a jet algorithm with a given radius parameter
   fastjet::JetDefinition jet_def(jetAlgorithm, R);
   fastjet::ClusterSequence clust_seq(input_particles, jet_def);

   // get the resulting jets ordered in pt
   auto inclusive_jets = sorted_by_pt(clust_seq.inclusive_jets(ptMin));
 
   auto hs = new THStack("hs","Jets;rapidity;#phi");
   int i=0;
   for (auto&& jet : inclusive_jets) {
     // get the constituents of the jet
     auto constituents = jet.constituents();
     auto h2Name = TString::Format("JetHist_%d",i);
     if (auto oldH2 = (TH2F*) gDirectory->GetObjectChecked(h2Name,"TH2F")) {
        delete oldH2;
   }
   auto h2 = new TH2F(h2Name,"JetHist",48, -5, 5, 48, 0, TMath::TwoPi());
   h2->SetFillColor(i+1);
   h2->SetLineWidth(1);
   h2->SetLineColor(kBlack);
   for (auto&& constituent : constituents){
      h2->Fill(constituent.rap(), constituent.phi_02pi(), constituent.perp());
   }
   hs->Add(h2);
   i++;
  }
  hs->Draw();// This creates the axes
  hs->GetXaxis()->SetTitleOffset(1.7);
  hs->GetYaxis()->SetTitleOffset(1.7);
  return hs;
}

%%python
inputFileName = 'Pythia-Z2jets-lhc-pileup-1ev.dat'
import os
if not os.path.exists(inputFileName):
    import urllib2
    response = urllib2.urlopen('https://raw.githubusercontent.com/dpiparo/swanExamples/master/notebooks/Pythia-Z2jets-lhc-pileup-1ev.dat')
    filecontent = response.read()   
    with open(inputFileName,"w") as f_out:
        f_out.write(filecontent)

TCanvas c;
c.SetTheta(90);
c.SetPhi(37);
unique_ptr<THStack> jetComponents(nullptr);

vector<fastjet::PseudoJet> input_particles;
fillInputParticles(input_particles, "Pythia-Z2jets-lhc-pileup-1ev.dat");

jetComponents.reset(getJetsComponents(input_particles, fastjet::antikt_algorithm, .8, 14.));
jetComponents->Draw("LEGO1 0");
c.Draw();