//----------------------------------------------------------------------
// Example on how to use this contribution
//
// run it with
//  ./example_npssub < ../data/Pythia-Zp2jets-lhc-pileup-1ev.dat
//----------------------------------------------------------------------

// $Id: example.cc 3001 2013-01-29 10:41:40Z soyez $
//
// Copyright (c) 2012-, Matteo Cacciari, Jihun Kim, Gavin P. Salam and Gregory Soyez
//
//----------------------------------------------------------------------
// This file is part of FastJet contrib.
//
// It is free software; you can redistribute it and/or modify it under
// the terms of the GNU General Public License as published by the
// Free Software Foundation; either version 2 of the License, or (at
// your option) any later version.
//
// It is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
// or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
// License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this code. If not, see <http://www.gnu.org/licenses/>.
//----------------------------------------------------------------------

//----------------------------------------------------------
// GS notes for CHS events
//----------------------------------------------------------

// Is we want to show an example of what happens for CHS events. we
// need to introduce a UserInfo class similar to what is in example 09
// in FastJet. For area wsubtraction, we'd just discard charged-PU
// tracks and that's mostly it (watching out that the Selector needs
// to be passed to the subtractor to properly handle the safety
// tests). But for NpC, we'd need a transformer (unless we use a loop)
// that rescales aprticles.
//
// Question: do we do that at all? Only Area? In this example or in
// another?

#include <iostream>
#include <iomanip>

#include "fastjet/PseudoJet.hh"
#include "fastjet/ClusterSequenceArea.hh"
#include "fastjet/Selector.hh"
#include "fastjet/tools/JetMedianBackgroundEstimator.hh"

//#include "fastjet/tools/Subtractor.hh"
#include "fastjet/Selector.hh"
#include "SafeSubtractor.hh"  // we'll use the implementation provided
                              // in this contrib
#include "GenericSubtractor.hh"

using namespace std;
using namespace fastjet;

void read_event(vector<PseudoJet> &hard_event, vector<PseudoJet> &full_event);
void print_jet(const PseudoJet &jet);

//----------------------------------------------------------------------
// User info associated with each PseudoJet that keeps tracks of the
// vertex number and cahrged info
//
// See e.g. example 09-user_info in FastJet
class MyUserInfo : public PseudoJet::UserInfoBase{
public:
  // default ctor
  //  - cahrge        the cahrge of the particle
  //  - vertex_number the id of the vertex it originates from
  MyUserInfo(const int & charge_in, const int & vertex_number_in) :
    _charge(charge_in), _vertex_number(vertex_number_in){}

  /// access to the PDG id
  int charge() const { return _charge;}
  
  /// access to the vertex number
  int vertex_number() const { return _vertex_number;}
  
protected:
  int _charge;         // the associated charge
  int _vertex_number;  // the associated vertex number
};

//----------------------------------------------------------------------
// a set of two selectors which allow to select (i) charged particles
// and (ii) particles coming from a given vertex

// worker to select charged particles
class SW_IsCharged : public SelectorWorker {
public:
  /// keep only charged particles
  ///
  /// Note that particles with no, or incompaticle user info will not pass
  virtual bool pass(const PseudoJet &p) const {
    return p.has_user_info<MyUserInfo>() 
        && p.user_info<MyUserInfo>().charge();
  }
};

// returns a selector from the previous worker
Selector SelectorIsCharged() {
  return Selector(new SW_IsCharged);
}

// worker to select particles from the 0th vertex
class SW_IsLeadingVertex : public SelectorWorker {
public:
  /// keep only particles from the 0th vertex
  ///
  /// Note that particles with no, or incompaticle user info will not pass
  virtual bool pass(const PseudoJet &p) const {
    return p.has_user_info<MyUserInfo>()
       && (p.user_info<MyUserInfo>().vertex_number()==0);
  }
};

// returns a selector from the previous worker
Selector SelectorIsLeadingVertex() {
  return Selector(new SW_IsLeadingVertex);
}

//----------------------------------------------------------------------
int main(){
  // read in input particles
  //
  // since we use here simulated data we can split the hard event
  // from the full (i.e. with pileup added) one
  //
  // (see also example 07 in FastJet)
  //----------------------------------------------------------
  vector<PseudoJet> hard_event, full_event;
  read_event(hard_event, full_event);

  // keep the particles up to 4 units in rapidity
  hard_event = SelectorAbsRapMax(4.0)(hard_event);
  full_event = SelectorAbsRapMax(4.0)(full_event);
  
  // do the clustering and get the jets
  //----------------------------------------------------------
  JetDefinition jet_def(antikt_algorithm, 0.7);
  AreaDefinition area_def(active_area_explicit_ghosts,
                          GhostedAreaSpec(SelectorAbsRapMax(4.0)));

  ClusterSequenceArea clust_seq_hard(hard_event, jet_def, area_def);
  ClusterSequenceArea clust_seq_full(full_event, jet_def, area_def);

  Selector sel_jets = SelectorNHardest(2) * SelectorAbsRapMax(3.0);

  vector<PseudoJet> hard_jets = sel_jets(clust_seq_hard.inclusive_jets());
  vector<PseudoJet> full_jets = sel_jets(clust_seq_full.inclusive_jets());

  // print the results without subtraction
  //----------------------------------------------------------
  cout << setprecision(4);
  cout << "# original hard jets" << endl;
  for (unsigned int i=0; i<hard_jets.size(); i++){
    const PseudoJet &jet = hard_jets[i];
    print_jet(jet);
  }
  cout << endl;

  cout << "# unsubtracted full jets" << endl;
  for (unsigned int i=0; i<full_jets.size(); i++){
    const PseudoJet &jet = full_jets[i];
    print_jet(jet);
  }
  cout << endl;


  // Neutral-proportional-to-charged subtraction
  //----------------------------------------------------------
  // Here, we just assume that the averaged fraction of cahrged tracks
  // in a patch is 0.612.
  //
  // Note: if we work with a CHS event, we should instead be keeping
  // the charged particles from PU vertices as ghosts, i.e. resale
  // them down by a factr epsilon<<1. In ythat ccse, we'd use
  //   charged_fraction = epsilon gamma/(1-gamma+epsilon gamma)
  // with gamma=0.612.
  //
  // Keeping of the PU charged tracks as ghosts is not necessary for
  // the area-median suibtraction
  const double charged_fraction = 0.612;
  
  // declare a NpC subtractor from this
  contrib::SafeNpCSubtractor npc_subtractor(charged_fraction,
					    SelectorIsCharged(),
					    SelectorIsLeadingVertex());

  cout << "# " << npc_subtractor.description() << endl;
  cout << "# NpC subtracted jets" << endl;
  for (unsigned int i=0; i<full_jets.size(); i++){
    PseudoJet subtracted_jet = npc_subtractor(full_jets[i]);
    print_jet(subtracted_jet);
  }
  cout << endl;

  return 0;
}

//------------------------------------------------------------------------
// read the event with and without pileup
void read_event(vector<PseudoJet> &hard_event, vector<PseudoJet> &full_event){
  string line;
  int  nsub  = 0; // counter to keep track of which sub-event we're reading
  while (getline(cin, line)) {
    istringstream linestream(line);
    // take substrings to avoid problems when there are extra "pollution"
    // characters (e.g. line-feed).
    if (line.substr(0,4) == "#END") {break;}
    if (line.substr(0,9) == "#SUBSTART") {
      // if more sub events follow, make copy of first one (the hard one) here
      if (nsub == 1) hard_event = full_event;
      nsub += 1;
    }
    if (line.substr(0,1) == "#") {continue;}
    double px,py,pz,E;
    int id, charge;
    linestream >> px >> py >> pz >> E >> id >> charge;
    PseudoJet particle(px,py,pz,E);

    // associate the user info to the particle
    particle.set_user_info(new MyUserInfo(charge, nsub-1));

    // push event onto back of full_event vector
    full_event.push_back(particle);
  }

  // if we have read in only one event, copy it across here...
  if (nsub == 1) hard_event = full_event;

  // if there was nothing in the event 
  if (nsub == 0) {
    cerr << "Error: read empty event\n";
    exit(-1);
  }

  cout << "# " << nsub-1 << " pileup events on top of the hard event" << endl;
}

//------------------------------------------------------------------------
// print iut some info about a giben jet
void print_jet(const PseudoJet &jet){
  cout << "pt = " << jet.pt()
       << ", rap = " << jet.rap()
       << ", mass = " << jet.m() << endl;
}