GenHFHadronMatcher.cc

Go to the documentation of this file.

// -*- C++ -*-
//
// Package:    GenHFHadronMatcher
// Class:      GenHFHadronMatcher
//
 
// Original Author:  Nazar Bartosik,DESY
 
// system include files
#include <memory>
#include <utility>
#include <vector>
#include <algorithm>
 
// user include files
#include "FWCore/Framework/interface/Frameworkfwd.h"
#include "FWCore/Framework/interface/global/EDProducer.h"
 
#include "FWCore/Framework/interface/Event.h"
#include "FWCore/Framework/interface/MakerMacros.h"
 
#include "FWCore/ParameterSet/interface/ParameterSet.h"
 
#include "FWCore/MessageLogger/interface/MessageLogger.h"
#include "FWCore/Framework/interface/ESHandle.h"
#include "FWCore/Framework/interface/EventSetup.h"
#include "FWCore/Utilities/interface/EDPutToken.h"
 
#include "DataFormats/HepMCCandidate/interface/GenParticle.h"
 
#include "DataFormats/JetReco/interface/GenJet.h"
#include "DataFormats/JetReco/interface/Jet.h"
#include "DataFormats/Math/interface/deltaR.h"
 
#include "DataFormats/JetMatching/interface/JetFlavourInfoMatching.h"
 
//
// class declaration
//
 
class GenHFHadronMatcher : public edm::global::EDProducer<> {
public:
  explicit GenHFHadronMatcher(const edm::ParameterSet &);
  ~GenHFHadronMatcher() override;
 
  static void fillDescriptions(edm::ConfigurationDescriptions &descriptions);
 
private:
  void produce(edm::StreamID, edm::Event &, const edm::EventSetup &) const override;
 
  std::vector<int> findHadronJets(const reco::GenParticleCollection *genParticles,
                                  const reco::JetFlavourInfoMatchingCollection *jetFlavourInfos,
                                  std::vector<int> &hadIndex,
                                  std::vector<reco::GenParticle> &hadMothersGenPart,
                                  std::vector<std::vector<int>> &hadMothersIndices,
                                  std::vector<int> &hadLeptonIndex,
                                  std::vector<int> &hadLeptonHadIndex,
                                  std::vector<int> &hadLeptonViaTau,
                                  std::vector<int> &hadFlavour,
                                  std::vector<int> &hadFromTopWeakDecay,
                                  std::vector<int> &hadBHadronId) const;
  int analyzeMothers(const reco::Candidate *thisParticle,
                     int &topDaughterQId,
                     int &topBarDaughterQId,
                     std::vector<const reco::Candidate *> &hadMothers,
                     std::vector<std::vector<int>> &hadMothersIndices,
                     std::set<const reco::Candidate *> *analyzedParticles,
                     const int prevPartIndex) const;
  bool putMotherIndex(std::vector<std::vector<int>> &hadMothersIndices, int partIndex, int mothIndex) const;
  bool isHadron(const int flavour, const reco::Candidate *thisParticle) const;
  bool isHadronPdgId(const int flavour, const int pdgId) const;
  bool isMesonPdgId(const int flavour, const int pdgId) const;
  bool isBaryonPdgId(const int flavour, const int pdgId) const;
  int flavourSign(const int pdgId) const;
  bool hasHadronDaughter(const int flavour, const reco::Candidate *thisParticle) const;
  int idInList(std::vector<const reco::Candidate *> particleList, const reco::Candidate *particle) const;
  int idInList(std::vector<int> list, const int value) const;
  int findInMothers(int idx,
                    std::vector<int> &mothChains,
                    const std::vector<std::vector<int>> &hadMothersIndices,
                    const std::vector<reco::GenParticle> &hadMothers,
                    int status,
                    int pdgId,
                    bool pdgAbs,
                    int stopId,
                    int firstLast,
                    bool verbose) const;
  bool isNeutralPdg(int pdgId) const;
 
  bool checkForLoop(std::vector<const reco::Candidate *> &particleChain, const reco::Candidate *particle) const;
  std::string getParticleName(int id) const;
 
  bool fixExtraSameFlavours(const unsigned int hadId,
                            const std::vector<int> &hadIndices,
                            const std::vector<reco::GenParticle> &hadMothers,
                            const std::vector<std::vector<int>> &hadMothersIndices,
                            const std::vector<int> &isFromTopWeakDecay,
                            const std::vector<std::vector<int>> &LastQuarkIds,
                            const std::vector<std::vector<int>> &LastQuarkMotherIds,
                            std::vector<int> &lastQuarkIndices,
                            std::vector<int> &hadronFlavour,
                            std::set<int> &checkedHadronIds,
                            const int lastQuarkIndex) const;
 
  // ----------member data ---------------------------
  const edm::EDGetTokenT<reco::GenParticleCollection> genParticlesToken_;
  const edm::EDGetTokenT<reco::JetFlavourInfoMatchingCollection> jetFlavourInfosToken_;
  const int flavour_;
  const bool noBBbarResonances_;
  const bool onlyJetClusteredHadrons_;
 
  const std::string flavourStr_;
  const edm::EDPutTokenT<std::vector<reco::GenParticle>> plusMothersToken_;
  const edm::EDPutTokenT<std::vector<std::vector<int>>> plusMothersIndicesToken_;
  const edm::EDPutTokenT<std::vector<int>> indexToken_;
  const edm::EDPutTokenT<std::vector<int>> flavourToken_;
  const edm::EDPutTokenT<std::vector<int>> jetIndexToken_;
  const edm::EDPutTokenT<std::vector<int>> leptonIndexToken_;
  const edm::EDPutTokenT<std::vector<int>> leptonHadronIndexToken_;
  const edm::EDPutTokenT<std::vector<int>> leptonViaTauToken_;
  const edm::EDPutTokenT<std::vector<int>> fromTopWeakDecayToken_;
  const edm::EDPutTokenT<std::vector<int>> bHadronIdToken_;
};
 
//
// constants, enums and typedefs
//
 
//
// static data member definitions
//
 
//
// constructors and destructor
//
namespace {
  std::string flavourName(int flavour) {
    if (flavour == 5) {
      return "B";
    } else if (flavour == 4) {
      return "C";
    }
    edm::LogError("GenHFHadronMatcher") << "Flavour option must be 4 (c-jet) or 5 (b-jet), but is: " << flavour
                                        << ". Correct this!";
    return std::string();
  }
}  // namespace
 
GenHFHadronMatcher::GenHFHadronMatcher(const edm::ParameterSet &cfg)
    : genParticlesToken_(consumes<reco::GenParticleCollection>(cfg.getParameter<edm::InputTag>("genParticles"))),
      jetFlavourInfosToken_(
          consumes<reco::JetFlavourInfoMatchingCollection>(cfg.getParameter<edm::InputTag>("jetFlavourInfos"))),
      flavour_{std::abs(cfg.getParameter<int>("flavour"))},
      noBBbarResonances_{cfg.getParameter<bool>("noBBbarResonances")},
      onlyJetClusteredHadrons_{cfg.getParameter<bool>("onlyJetClusteredHadrons")},
      flavourStr_{flavourName(flavour_)},
      plusMothersToken_{produces<std::vector<reco::GenParticle>>(
          "gen" + flavourStr_ +
          "HadPlusMothers")},  // All mothers in all decay chains above any hadron of specified flavour
      plusMothersIndicesToken_{produces<std::vector<std::vector<int>>>(
          "gen" + flavourStr_ + "HadPlusMothersIndices")},  // Indices of mothers of each hadMother
      indexToken_{
          produces<std::vector<int>>("gen" + flavourStr_ + "HadIndex")},  // Index of hadron in the vector of hadMothers
      flavourToken_{produces<std::vector<int>>(
          "gen" + flavourStr_ +
          "HadFlavour")},  // PdgId of the first non-b(c) quark mother with sign corresponding to hadron charge
      jetIndexToken_{produces<std::vector<int>>(
          "gen" + flavourStr_ + "HadJetIndex")},  // Index of genJet matched to each hadron by jet clustering algorithm
      leptonIndexToken_{produces<std::vector<int>>(
          "gen" + flavourStr_ +
          "HadLeptonIndex")},  // Index of lepton found among the hadron decay products in the list of mothers
      leptonHadronIndexToken_{produces<std::vector<int>>(
          "gen" + flavourStr_ + "HadLeptonHadronIndex")},  // Index of hadron the lepton is associated to
      leptonViaTauToken_{produces<std::vector<int>>(
          "gen" + flavourStr_ + "HadLeptonViaTau")},  // Whether lepton comes directly from hadron or via tau decay
      fromTopWeakDecayToken_{produces<std::vector<int>>(
          "gen" + flavourStr_ +
          "HadFromTopWeakDecay")},  // Tells whether the hadron appears in the chain after top decay
      bHadronIdToken_{produces<std::vector<int>>("gen" + flavourStr_ + "HadBHadronId")}
// Index of a b-hadron which the current hadron comes from (for c-hadrons)
{
  // Hadron matching products
}
 
GenHFHadronMatcher::~GenHFHadronMatcher() {}
 
// ------------ method fills 'descriptions' with the allowed parameters for the module  ------------
void GenHFHadronMatcher::fillDescriptions(edm::ConfigurationDescriptions &descriptions) {
  edm::ParameterSetDescription desc;
  desc.add<edm::InputTag>("genParticles")
      ->setComment("Collection of GenParticle objects which contains all particles produced in the event");
  desc.add<edm::InputTag>("jetFlavourInfos")
      ->setComment(
          "Output from the JetFlavour tool. Contains information about partons/hadrons/leptons associated to jets");
  desc.add<bool>("noBBbarResonances", true)->setComment("Whether resonances should not be treated as hadrons");
  desc.add<bool>("onlyJetClusteredHadrons", false)
      ->setComment("Whether only hadrons that are matched to jets should be analysed. Runs x1000 faster in Sherpa");
  desc.add<int>("flavour", 5)->setComment("Flavour of weakly decaying hadron that should be analysed (4-c, 5-b)");
  descriptions.add("matchGenHFHadron", desc);
}
 
//
// member functions
//
 
// ------------ method called to produce the data  ------------
void GenHFHadronMatcher::produce(edm::StreamID, edm::Event &evt, const edm::EventSetup &setup) const {
  using namespace edm;
 
  edm::Handle<reco::GenParticleCollection> genParticles;
  evt.getByToken(genParticlesToken_, genParticles);
 
  edm::Handle<reco::JetFlavourInfoMatchingCollection> jetFlavourInfos;
  evt.getByToken(jetFlavourInfosToken_, jetFlavourInfos);
 
  // Defining adron matching variables
  std::vector<reco::GenParticle> hadMothers;
  std::vector<std::vector<int>> hadMothersIndices;
  std::vector<int> hadIndex;
  std::vector<int> hadFlavour;
  std::vector<int> hadJetIndex;
  std::vector<int> hadLeptonIndex;
  std::vector<int> hadLeptonHadIndex;
  std::vector<int> hadLeptonViaTau;
  std::vector<int> hadFromTopWeakDecay;
  std::vector<int> hadBHadronId;
 
  hadJetIndex = findHadronJets(genParticles.product(),
                               jetFlavourInfos.product(),
                               hadIndex,
                               hadMothers,
                               hadMothersIndices,
                               hadLeptonIndex,
                               hadLeptonHadIndex,
                               hadLeptonViaTau,
                               hadFlavour,
                               hadFromTopWeakDecay,
                               hadBHadronId);
 
  // Putting products to the event
  evt.emplace(plusMothersToken_, std::move(hadMothers));
  evt.emplace(plusMothersIndicesToken_, std::move(hadMothersIndices));
  evt.emplace(indexToken_, std::move(hadIndex));
  evt.emplace(flavourToken_, std::move(hadFlavour));
  evt.emplace(jetIndexToken_, std::move(hadJetIndex));
  evt.emplace(leptonIndexToken_, std::move(hadLeptonIndex));
  evt.emplace(leptonHadronIndexToken_, std::move(hadLeptonHadIndex));
  evt.emplace(leptonViaTauToken_, std::move(hadLeptonViaTau));
  evt.emplace(fromTopWeakDecayToken_, std::move(hadFromTopWeakDecay));
  evt.emplace(bHadronIdToken_, std::move(hadBHadronId));
}
 
std::vector<int> GenHFHadronMatcher::findHadronJets(const reco::GenParticleCollection *genParticles,
                                                    const reco::JetFlavourInfoMatchingCollection *jetFlavourInfos,
                                                    std::vector<int> &hadIndex,
                                                    std::vector<reco::GenParticle> &hadMothers,
                                                    std::vector<std::vector<int>> &hadMothersIndices,
                                                    std::vector<int> &hadLeptonIndex,
                                                    std::vector<int> &hadLeptonHadIndex,
                                                    std::vector<int> &hadLeptonViaTau,
                                                    std::vector<int> &hadFlavour,
                                                    std::vector<int> &hadFromTopWeakDecay,
                                                    std::vector<int> &hadBHadronId) const {
  std::vector<int> hadJetIndex;
  std::vector<const reco::Candidate *> hadMothersCand;
 
  int topDaughterQId = -1;
  int topBarDaughterQId = -1;
 
  // Looping over all jets to get hadrons associated to them
  for (reco::JetFlavourInfoMatchingCollection::const_iterator i_info = jetFlavourInfos->begin();
       i_info != jetFlavourInfos->end();
       ++i_info) {
    reco::JetFlavourInfo jetInfo = i_info->second;
    const int jetIndex = i_info - jetFlavourInfos->begin();
    // Looping over each hadron associated with the jet and finding its origin
    const reco::GenParticleRefVector &hadronsInJet = flavour_ == 5   ? jetInfo.getbHadrons()
                                                     : flavour_ == 4 ? jetInfo.getcHadrons()
                                                                     : reco::GenParticleRefVector();
    for (reco::GenParticleRefVector::const_iterator hadron = hadronsInJet.begin(); hadron != hadronsInJet.end();
         ++hadron) {
      // Check that the hadron satisfies criteria configured in the module
      if (!isHadron(flavour_, (&**hadron)))
        continue;
      if (hasHadronDaughter(flavour_, (reco::Candidate *)(&**hadron)))
        continue;
      // Scanning the chain starting from the hadron
      int hadronIndex = analyzeMothers((reco::Candidate *)(&**hadron),
                                       topDaughterQId,
                                       topBarDaughterQId,
                                       hadMothersCand,
                                       hadMothersIndices,
                                       nullptr,
                                       -1);
      // Storing the index of the hadron to the list
      hadIndex.push_back(hadronIndex);
      hadJetIndex.push_back(jetIndex);  // Putting jet index to the result list
    }
  }  // End of loop over jets
 
  // Access all hadrons which are not associated with jets, if requested
  if (!onlyJetClusteredHadrons_) {
    for (reco::GenParticleCollection::const_iterator i_particle = genParticles->begin();
         i_particle != genParticles->end();
         ++i_particle) {
      const reco::GenParticle *thisParticle = &*i_particle;
      if (!isHadron(flavour_, thisParticle))
        continue;
      // Skipping the hadron if it was already found directly from jets
      if (std::find(hadMothersCand.begin(), hadMothersCand.end(), thisParticle) != hadMothersCand.end())
        continue;
 
      // Scanning the chain starting from the hadron
      int hadronIndex = analyzeMothers(
          thisParticle, topDaughterQId, topBarDaughterQId, hadMothersCand, hadMothersIndices, nullptr, -1);
      // Storing the index of the hadron to the list
      hadIndex.push_back(hadronIndex);
      hadJetIndex.push_back(-1);  // Jet index undefined
    }
  }
 
  // Transfering Candidates to the list of processed particles for further analysis
  for (int i = 0; i < (int)hadMothersCand.size(); i++) {
    const reco::GenParticle *particle = dynamic_cast<const reco::GenParticle *>(hadMothersCand.at(i));
    hadMothers.push_back(*particle);
  }
 
  // Adding leptons from hadron decays
  for (reco::GenParticleCollection::const_iterator i_particle = genParticles->begin();
       i_particle != genParticles->end();
       ++i_particle) {
    const reco::GenParticle lepton = *i_particle;
    const int pdg_abs = lepton.pdgId();
    // Skipping if not a lepton: e/mu
    if (pdg_abs != 11 && pdg_abs != 13)
      continue;
    bool leptonViaTau = false;
    const reco::Candidate *leptonMother = lepton.mother();
    if (!leptonMother)
      continue;
    // Taking next mother if direct mother is a tau
    if (std::abs(leptonMother->pdgId()) == 15) {
      leptonViaTau = true;
      leptonMother = leptonMother->mother();
    }
    // Skipping this lepton if its mother is not a proper hadron
    if (leptonMother == nullptr or !isHadron(flavour_, leptonMother))
      continue;
    // Finding the index of this hadron in the list of analysed particles
    size_t leptonHadronParticleIndex =
        std::find(hadMothersCand.begin(), hadMothersCand.end(), leptonMother) - hadMothersCand.begin();
    if (leptonHadronParticleIndex >= hadMothersCand.size())
      continue;
    // Finding the actual hadron index among those that were found
    size_t leptonHadronIndex =
        std::find(hadIndex.begin(), hadIndex.end(), leptonHadronParticleIndex) - hadIndex.begin();
    if (leptonHadronIndex >= hadIndex.size())
      continue;
    // Putting the lepton, its index and hadron index to the corresponding lists
    hadMothers.push_back(lepton);
    const int leptonIndex = hadMothersCand.size() - 1;
    hadLeptonIndex.push_back(leptonIndex);
    hadLeptonViaTau.push_back(leptonViaTau);
    hadLeptonHadIndex.push_back(leptonHadronIndex);
  }
 
  // Checking mothers of hadrons in order to assign flags (where the hadron comes from)
  unsigned int nHad = hadIndex.size();
 
  std::vector<std::vector<int>> LastQuarkMotherIds;
  std::vector<std::vector<int>> LastQuarkIds;
  std::vector<int> lastQuarkIndices(nHad, -1);
 
  // Looping over all hadrons
  for (unsigned int hadNum = 0; hadNum < nHad; hadNum++) {
    unsigned int hadIdx = hadIndex.at(hadNum);  // Index of hadron in the hadMothers
 
    std::vector<int> FirstQuarkId;
    std::vector<int> LastQuarkId;
    std::vector<int> LastQuarkMotherId;
 
    const int hadronFlavourSign = flavourSign(hadMothers.at(hadIdx).pdgId());
 
    // Searching only first quark in the chain with the same flavour as hadron
    if (hadronFlavourSign != 0) {
      findInMothers(
          hadIdx, FirstQuarkId, hadMothersIndices, hadMothers, 0, hadronFlavourSign * flavour_, false, -1, 1, false);
    }
    // Searching for quarks with both flavours since it is a bb/cc resonance
    else {
      findInMothers(hadIdx, FirstQuarkId, hadMothersIndices, hadMothers, 0, flavour_, false, -1, 1, false);
      findInMothers(hadIdx, FirstQuarkId, hadMothersIndices, hadMothers, 0, -1 * flavour_, false, -1, 1, false);
    }
 
    // Finding last quark for each first quark
    for (unsigned int qId = 0; qId < FirstQuarkId.size(); qId++) {
      // Identifying the flavour of the first quark to find the last quark of the same flavour
      const int quarkFlavourSign = flavourSign(hadMothers.at(FirstQuarkId.at(qId)).pdgId());
      // Finding last quark of the hadron starting from the first quark
      findInMothers(FirstQuarkId.at(qId),
                    LastQuarkId,
                    hadMothersIndices,
                    hadMothers,
                    0,
                    quarkFlavourSign * flavour_,
                    false,
                    -1,
                    2,
                    false);
    }  // End of loop over all first quarks of the hadron
 
    // Setting initial flavour of the hadron
    int hadronFlavour = 0;
 
    // Initialising pairs of last quark index and its distance from the hadron (to sort by distance)
    std::vector<std::pair<double, int>> lastQuark_dR_id_pairs;
 
    // Finding the closest quark in dR
    for (unsigned int qId = 0; qId < LastQuarkId.size(); qId++) {
      int qIdx = LastQuarkId.at(qId);
      // Calculating the dR between hadron and quark
      float dR = deltaR(hadMothers.at(hadIdx).eta(),
                        hadMothers.at(hadIdx).phi(),
                        hadMothers.at(qIdx).eta(),
                        hadMothers.at(qIdx).phi());
 
      std::pair<double, int> dR_hadId_pair(dR, qIdx);
      lastQuark_dR_id_pairs.push_back(dR_hadId_pair);
    }  // End of loop over all last quarks of the hadron
 
    std::sort(lastQuark_dR_id_pairs.begin(), lastQuark_dR_id_pairs.end());
 
    if (lastQuark_dR_id_pairs.size() > 1) {
      double dRratio =
          (lastQuark_dR_id_pairs.at(1).first - lastQuark_dR_id_pairs.at(0).first) / lastQuark_dR_id_pairs.at(1).first;
      int qIdx_closest = lastQuark_dR_id_pairs.at(0).second;
      LastQuarkId.clear();
      if (dRratio > 0.5)
        LastQuarkId.push_back(qIdx_closest);
      else
        for (std::pair<double, int> qIdDrPair : lastQuark_dR_id_pairs)
          LastQuarkId.push_back(qIdDrPair.second);
    }
    for (int qIdx : LastQuarkId) {
      int qmIdx = hadMothersIndices.at(qIdx).at(0);
      LastQuarkMotherId.push_back(qmIdx);
    }
 
    if ((int)LastQuarkId.size() > 0)
      lastQuarkIndices.at(hadNum) = 0;  // Setting the first quark in array as a candidate if it exists
 
    LastQuarkIds.push_back(LastQuarkId);
 
    LastQuarkMotherIds.push_back(LastQuarkMotherId);
 
    if (LastQuarkMotherId.empty()) {
      hadronFlavour = 0;
    } else {
      int qIdx = LastQuarkId.at(lastQuarkIndices.at(hadNum));
      int qFlav = flavourSign(hadMothers.at(qIdx).pdgId());
      hadronFlavour = qFlav * std::abs(hadMothers.at(LastQuarkMotherId.at(lastQuarkIndices.at(hadNum))).pdgId());
    }
    hadFlavour.push_back(hadronFlavour);  // Adding hadron flavour to the list of flavours
 
    // Checking whether hadron comes from the Top weak decay
    int isFromTopWeakDecay = 1;
    std::vector<int> checkedParticles;
    if (hadFlavour.at(hadNum) != 0) {
      int lastQIndex = LastQuarkId.at(lastQuarkIndices.at(hadNum));
      bool fromTB = topDaughterQId >= 0 ? findInMothers(lastQIndex,
                                                        checkedParticles,
                                                        hadMothersIndices,
                                                        hadMothers,
                                                        -1,
                                                        0,
                                                        false,
                                                        topDaughterQId,
                                                        2,
                                                        false) >= 0
                                        : false;
      checkedParticles.clear();
      bool fromTbarB = topBarDaughterQId >= 0 ? findInMothers(lastQIndex,
                                                              checkedParticles,
                                                              hadMothersIndices,
                                                              hadMothers,
                                                              -1,
                                                              0,
                                                              false,
                                                              topBarDaughterQId,
                                                              2,
                                                              false) >= 0
                                              : false;
      checkedParticles.clear();
      if (!fromTB && !fromTbarB) {
        isFromTopWeakDecay = 0;
      }
    } else
      isFromTopWeakDecay = 2;
    hadFromTopWeakDecay.push_back(isFromTopWeakDecay);
    int bHadronMotherId =
        findInMothers(hadIdx, checkedParticles, hadMothersIndices, hadMothers, 0, 555555, true, -1, 1, false);
    hadBHadronId.push_back(bHadronMotherId);
 
    if (!LastQuarkMotherId.empty()) {
      std::set<int> checkedHadronIds;
      fixExtraSameFlavours(hadNum,
                           hadIndex,
                           hadMothers,
                           hadMothersIndices,
                           hadFromTopWeakDecay,
                           LastQuarkIds,
                           LastQuarkMotherIds,
                           lastQuarkIndices,
                           hadFlavour,
                           checkedHadronIds,
                           0);
    }
 
  }  // End of loop over all hadrons
 
  return hadJetIndex;
}
 
int GenHFHadronMatcher::idInList(std::vector<const reco::Candidate *> particleList,
                                 const reco::Candidate *particle) const {
  const unsigned int position = std::find(particleList.begin(), particleList.end(), particle) - particleList.begin();
  if (position >= particleList.size())
    return -1;
 
  return position;
}
 
int GenHFHadronMatcher::idInList(std::vector<int> list, const int value) const {
  const unsigned int position = std::find(list.begin(), list.end(), value) - list.begin();
  if (position >= list.size())
    return -1;
 
  return position;
}
 
bool GenHFHadronMatcher::isHadron(const int flavour, const reco::Candidate *thisParticle) const {
  return isHadronPdgId(flavour, thisParticle->pdgId());
}
 
bool GenHFHadronMatcher::isHadronPdgId(const int flavour, const int pdgId) const {
  if (isBaryonPdgId(flavour, pdgId) || isMesonPdgId(flavour, pdgId))
    return true;
 
  return false;
}
 
bool GenHFHadronMatcher::isMesonPdgId(const int flavour, const int pdgId) const {
  const int flavour_abs = std::abs(flavour);
  if (flavour_abs != 5 && flavour_abs != 4)
    return false;
  const int pdgId_abs = std::abs(pdgId);
 
  if (pdgId_abs / 100 % 10 != flavour_abs)
    return false;
  // Excluding baryons
  if (pdgId_abs / 1000 == flavour_abs)
    return false;
  // Excluding bb/cc resonances if required
  if (noBBbarResonances_ && pdgId_abs / 10 % 100 == 11 * flavour_abs)
    return false;
 
  return true;
}
 
bool GenHFHadronMatcher::isBaryonPdgId(const int flavour, const int pdgId) const {
  const int flavour_abs = std::abs(flavour);
  if (flavour_abs != 5 && flavour_abs != 4)
    return false;
  const int pdgId_abs = std::abs(pdgId);
 
  if (pdgId_abs / 1000 != flavour_abs)
    return false;
 
  return true;
}
 
int GenHFHadronMatcher::flavourSign(const int pdgId) const {
  int flavourSign = pdgId / std::abs(pdgId);
  // B mesons have opposite sign
  if (isMesonPdgId(5, pdgId))
    flavourSign *= -1;
  // Returning 0 for bb/cc resonances
  if (pdgId % 1000 / 10 / 11 > 0)
    flavourSign = 0;
 
  return flavourSign;
}
 
bool GenHFHadronMatcher::hasHadronDaughter(const int flavour, const reco::Candidate *thisParticle) const {
  // Looping through daughters of the particle
  bool hasDaughter = false;
  for (int k = 0; k < (int)thisParticle->numberOfDaughters(); k++) {
    if (!isHadron(flavour, thisParticle->daughter(k))) {
      continue;
    }
    hasDaughter = true;
    break;
  }
 
  return hasDaughter;
}
 
int GenHFHadronMatcher::analyzeMothers(const reco::Candidate *thisParticle,
                                       int &topDaughterQId,
                                       int &topBarDaughterQId,
                                       std::vector<const reco::Candidate *> &hadMothers,
                                       std::vector<std::vector<int>> &hadMothersIndices,
                                       std::set<const reco::Candidate *> *analyzedParticles,
                                       const int prevPartIndex) const {
  // Getting the index of the particle which is a hadron in the first call
  int hadronIndex = -1;  // Index of the hadron that is returned by this function
  int index = idInList(hadMothers, thisParticle);
  if (index < 0) {  // If hadron is not in the list of mothers yet
    hadMothers.push_back(thisParticle);
    hadronIndex = hadMothers.size() - 1;
  } else {  // If hadron is in the list of mothers already
    hadronIndex = index;
  }
 
  int partIndex = -1;  // Index of particle being checked in the list of mothers
  partIndex = idInList(hadMothers, thisParticle);
 
  // Checking whether this particle is already in the chain of analyzed particles in order to identify a loop
  bool isLoop = false;
  if (!analyzedParticles) {
    analyzedParticles = new std::set<const reco::Candidate *>;
  }
  for (unsigned int i = 0; i < analyzedParticles->size(); i++) {
    if (analyzedParticles->count(thisParticle) <= 0) {
      continue;
    }
    isLoop = true;
    break;
  }
 
  // If a loop has been detected
  if (isLoop) {
    if (prevPartIndex >= 0) {
      putMotherIndex(hadMothersIndices, prevPartIndex, -1);  // Setting mother index of previous particle to -1
    }
    return hadronIndex;  // Stopping further processing of the current chain
  }
  analyzedParticles->insert(thisParticle);
 
  // Putting the mothers to the list of mothers
  for (size_t iMother = 0; iMother < thisParticle->numberOfMothers(); ++iMother) {
    const reco::Candidate *mother = thisParticle->mother(iMother);
    int mothIndex = idInList(hadMothers, mother);
    if (mothIndex == partIndex && partIndex >= 0) {
      continue;  // Skipping the mother that is its own daughter
    }
 
    // If this mother isn't yet in the list and hadron or lepton is in the list
    if (mothIndex < 0) {
      hadMothers.push_back(mother);
      mothIndex = hadMothers.size() - 1;
    }
    // If hadron has already been found in current chain and the mother isn't a duplicate of the particle being checked
    if (mothIndex != partIndex && partIndex >= 0) {
      putMotherIndex(hadMothersIndices, partIndex, mothIndex);  // Putting the index of mother for current particle
    }
    analyzeMothers(
        mother, topDaughterQId, topBarDaughterQId, hadMothers, hadMothersIndices, analyzedParticles, partIndex);
    // Setting the id of the particle which is a quark from the top decay
    if (std::abs(mother->pdgId()) == 6) {
      int &bId = mother->pdgId() < 0 ? topBarDaughterQId : topDaughterQId;
      int thisFlav = std::abs(thisParticle->pdgId());
      if (bId < 0) {
        if (thisFlav <= 5)
          bId = partIndex;
      } else {
        int bIdFlav = std::abs(hadMothers.at(bId)->pdgId());
        if (bIdFlav != 5 && thisFlav == 5)
          bId = partIndex;
        else if (thisFlav == 5 && thisParticle->pt() > hadMothers.at(bId)->pt())
          bId = partIndex;
      }  // If daughter quark of the top not found yet
    }    // If the mother is a top quark and hadron has been found
  }      // End of loop over mothers
 
  analyzedParticles->erase(thisParticle);  // Removing current particle from the current chain that is being analyzed
 
  if (partIndex < 0) {
    return hadronIndex;  // Safety check
  }
 
  // Adding -1 to the list of mother indices for current particle if it has no mothers (for consistency between numbering of indices and mothers)
  if ((int)thisParticle->numberOfMothers() <= 0) {
    putMotherIndex(hadMothersIndices, partIndex, -1);
  }
 
  return hadronIndex;
}
 
bool GenHFHadronMatcher::putMotherIndex(std::vector<std::vector<int>> &hadMothersIndices,
                                        int partIndex,
                                        int mothIndex) const {
  // Putting vector of mothers indices for the given particle
  bool inList = false;
  if (partIndex < 0) {
    return false;
  }
 
  while ((int)hadMothersIndices.size() <= partIndex) {  // If there is no list of mothers for current particle yet
    std::vector<int> mothersIndices;
    hadMothersIndices.push_back(mothersIndices);
  }
 
  std::vector<int> *hadMotherIndices = &hadMothersIndices.at(partIndex);
  // Removing other mothers if particle must have no mothers
  if (mothIndex == -1) {
    hadMotherIndices->clear();
  } else {
    // Checking if current mother is already in the list of theParticle's mothers
    for (int k = 0; k < (int)hadMotherIndices->size(); k++) {
      if (hadMotherIndices->at(k) != mothIndex && hadMotherIndices->at(k) != -1) {
        continue;
      }
      inList = true;
      break;
    }
  }
  // Adding current mother to the list of mothers of this particle
  if (!inList) {
    hadMotherIndices->push_back(mothIndex);
  }
 
  return inList;
}
 
int GenHFHadronMatcher::findInMothers(int idx,
                                      std::vector<int> &mothChains,
                                      const std::vector<std::vector<int>> &hadMothersIndices,
                                      const std::vector<reco::GenParticle> &hadMothers,
                                      int status,
                                      int pdgId,
                                      bool pdgAbs = false,
                                      int stopId = -1,
                                      int firstLast = 0,
                                      bool verbose = false) const {
  int foundStopId = -1;
  int pdg_1 = hadMothers.at(idx).pdgId();
 
  if ((int)hadMothersIndices.size() <= idx) {
    if (verbose) {
      printf(" Stopping checking particle %d. No mothers are stored.\n", idx);
    }
    return -1;  // Skipping if no mothers are stored for this particle
  }
 
  if (std::abs(hadMothers.at(idx).pdgId()) > 10 && std::abs(hadMothers.at(idx).pdgId()) < 19)
    printf("Lepton: %d\n", hadMothers.at(idx).pdgId());
 
  std::vector<int> mothers = hadMothersIndices.at(idx);
  unsigned int nMothers = mothers.size();
  bool isCorrect = false;  // Whether current particle is what is being searched
  if (verbose) {
    if (std::abs(hadMothers.at(idx).pdgId()) == 2212) {
      printf("Chk:  %d\tpdg: %d\tstatus: %d", idx, hadMothers.at(idx).pdgId(), hadMothers.at(idx).status());
    } else {
      printf(" Chk:  %d(%d mothers)\tpdg: %d\tstatus: %d\tPt: %.3f\tEta: %.3f",
             idx,
             nMothers,
             hadMothers.at(idx).pdgId(),
             hadMothers.at(idx).status(),
             hadMothers.at(idx).pt(),
             hadMothers.at(idx).eta());
    }
  }
  bool hasCorrectMothers = true;
  if (nMothers < 1)
    hasCorrectMothers = false;
  else if (mothers.at(0) < 0)
    hasCorrectMothers = false;
  if (!hasCorrectMothers) {
    if (verbose)
      printf("    NO CORRECT MOTHER\n");
    return -1;
  }
  // Stopping if the particular particle has been found
  if (stopId >= 0 && idx == stopId)
    return idx;
 
  // Stopping if the hadron of particular flavour has been found
  if (pdgId % 111111 == 0 && pdgId != 0) {
    if (isHadronPdgId(pdgId / 111111, hadMothers.at(idx).pdgId())) {
      return idx;
    }
  }
 
  // Checking whether current mother satisfies selection criteria
  if (((hadMothers.at(idx).pdgId() == pdgId && pdgAbs == false) ||
       (std::abs(hadMothers.at(idx).pdgId()) == std::abs(pdgId) && pdgAbs == true)) &&
      (hadMothers.at(idx).status() == status || status == 0) && hasCorrectMothers) {
    isCorrect = true;
    // Adding to the list of candidates if not there and if mother of this quark has correct flavour sign
    const bool inList = std::find(mothChains.begin(), mothChains.end(), idx) != mothChains.end();
    if (!inList && mothers.at(0) >= 0 && (hadMothers.at(idx).pdgId() * pdgId > 0 || !pdgAbs)) {
      if (firstLast == 0 || firstLast == 1) {
        mothChains.push_back(idx);
      }
      if (verbose) {
        printf("   *");
      }
    }
    if (verbose) {
      printf("   +++");
    }
  }
  if (verbose) {
    printf("\n");
  }
 
  if (isCorrect && firstLast == 1) {
    return -1;  // Stopping if only the first particle in the chain is looked for
  }
 
  // Checking next level mothers
  unsigned int nDifferingMothers = 0;
  for (unsigned int i = 0; i < nMothers; i++) {
    int idx2 = mothers.at(i);
    if (idx2 < 0) {
      if (verbose)
        printf("^^^ Has no mother\n");
      continue;  // Skipping if mother's id is -1 (no mother), that means current particle is a proton
    }
    if (idx2 == idx) {
      if (verbose)
        printf("^^^ Stored as its own mother\n");
      continue;  // Skipping if particle is stored as its own mother
    }
    int pdg_2 = hadMothers[idx2].pdgId();
    // Inverting the flavour if bb oscillation detected
    if (isHadronPdgId(pdgId, pdg_1) && isHadronPdgId(pdgId, pdg_2) && pdg_1 * pdg_2 < 0) {
      pdgId *= -1;
      if (verbose)
        printf("######### Inverting flavour of the hadron\n");
    }
    // Counting how many mothers are different from this particle
    if ((std::abs(pdg_2) != std::abs(pdgId) && pdgAbs == true) || (pdg_2 != pdgId && pdgAbs == false)) {
      nDifferingMothers++;
    }
 
    // Checking next level mother
    if (verbose) {
      printf("Checking mother %d out of %d mothers (%d -> %d), looking for pdgId: %d\n", i, nMothers, idx, idx2, pdgId);
    }
    if (firstLast == 2 && pdg_1 != pdg_2)
      continue;  // Requiring the same flavour when finding the last quark
    foundStopId = findInMothers(
        idx2, mothChains, hadMothersIndices, hadMothers, status, pdgId, pdgAbs, stopId, firstLast, verbose);
  }
  // Storing this particle if all its mothers are of another type and the last of its kind should be stored
  if (firstLast == 2 && isCorrect && nDifferingMothers >= nMothers) {
    if (verbose) {
      printf("Checking particle %d once more to store it as the last quark\n", idx);
    }
    foundStopId = findInMothers(idx, mothChains, hadMothersIndices, hadMothers, 0, pdgId, pdgAbs, stopId, 1, verbose);
  }
 
  return foundStopId;
}
 
bool GenHFHadronMatcher::isNeutralPdg(int pdgId) const {
  const int neutralPdgs_array[] = {9, 21, 22, 23, 25};
  const std::vector<int> neutralPdgs(neutralPdgs_array, neutralPdgs_array + sizeof(neutralPdgs_array) / sizeof(int));
  if (std::find(neutralPdgs.begin(), neutralPdgs.end(), std::abs(pdgId)) == neutralPdgs.end())
    return false;
 
  return true;
}
 
bool GenHFHadronMatcher::fixExtraSameFlavours(const unsigned int hadId,
                                              const std::vector<int> &hadIndices,
                                              const std::vector<reco::GenParticle> &hadMothers,
                                              const std::vector<std::vector<int>> &hadMothersIndices,
                                              const std::vector<int> &isFromTopWeakDecay,
                                              const std::vector<std::vector<int>> &LastQuarkIds,
                                              const std::vector<std::vector<int>> &LastQuarkMotherIds,
                                              std::vector<int> &lastQuarkIndices,
                                              std::vector<int> &hadronFlavour,
                                              std::set<int> &checkedHadronIds,
                                              const int lastQuarkIndex) const {
  if (checkedHadronIds.count(hadId) != 0)
    return false;                  // Hadron already checked previously and should be skipped
  checkedHadronIds.insert(hadId);  // Putting hadron to the list of checked ones in this run
 
  if (lastQuarkIndex < 0)
    return false;
  if ((int)LastQuarkIds.at(hadId).size() < lastQuarkIndex + 1)
    return false;
  int LastQuarkId = LastQuarkIds.at(hadId).at(lastQuarkIndex);
  int LastQuarkMotherId = LastQuarkMotherIds.at(hadId).at(lastQuarkIndex);
  int qmFlav = hadMothers.at(LastQuarkId).pdgId() < 0 ? -1 : 1;
  int hadFlavour = qmFlav * std::abs(hadMothers.at(LastQuarkMotherId).pdgId());
  bool ambiguityResolved = true;
  // If last quark has inconsistent flavour with its mother, setting the hadron flavour to gluon
  if ((hadMothers.at(LastQuarkId).pdgId() * hadMothers.at(LastQuarkMotherId).pdgId() < 0 &&
       !isNeutralPdg(hadMothers.at(LastQuarkMotherId).pdgId())) ||
      // If particle not coming from the Top weak decay has Top flavour
      (std::abs(hadronFlavour.at(hadId)) == 6 && isFromTopWeakDecay.at(hadId) == 0)) {
    if ((int)LastQuarkIds.at(hadId).size() > lastQuarkIndex + 1)
      fixExtraSameFlavours(hadId,
                           hadIndices,
                           hadMothers,
                           hadMothersIndices,
                           isFromTopWeakDecay,
                           LastQuarkIds,
                           LastQuarkMotherIds,
                           lastQuarkIndices,
                           hadronFlavour,
                           checkedHadronIds,
                           lastQuarkIndex + 1);
    else
      hadronFlavour.at(hadId) = qmFlav * 21;
    return true;
  }
 
  int nSameFlavourHadrons = 0;
  // Looping over all previous hadrons
  for (unsigned int iHad = 0; iHad < hadronFlavour.size(); iHad++) {
    if (iHad == hadId)
      continue;
    int theLastQuarkIndex = lastQuarkIndices.at(iHad);
    if (theLastQuarkIndex < 0)
      continue;
    if ((int)LastQuarkMotherIds.at(iHad).size() <= theLastQuarkIndex)
      continue;
    int theLastQuarkMotherId = LastQuarkMotherIds.at(iHad).at(theLastQuarkIndex);
    int theHadFlavour = hadronFlavour.at(iHad);
    // Skipping hadrons with different flavour
    if (theHadFlavour == 0 || std::abs(theHadFlavour) == 21)
      continue;
    if (theHadFlavour != hadFlavour || theLastQuarkMotherId != LastQuarkMotherId)
      continue;
    ambiguityResolved = false;
    nSameFlavourHadrons++;
 
    // Checking other b-quark mother candidates of this hadron
    if ((int)LastQuarkIds.at(hadId).size() > lastQuarkIndex + 1) {
      if (fixExtraSameFlavours(hadId,
                               hadIndices,
                               hadMothers,
                               hadMothersIndices,
                               isFromTopWeakDecay,
                               LastQuarkIds,
                               LastQuarkMotherIds,
                               lastQuarkIndices,
                               hadronFlavour,
                               checkedHadronIds,
                               lastQuarkIndex + 1)) {
        ambiguityResolved = true;
        break;
      }
    } else
        // Checking other b-quark mother candidates of the particular previous hadron
        if ((int)LastQuarkIds.at(iHad).size() > theLastQuarkIndex + 1) {
      if (fixExtraSameFlavours(iHad,
                               hadIndices,
                               hadMothers,
                               hadMothersIndices,
                               isFromTopWeakDecay,
                               LastQuarkIds,
                               LastQuarkMotherIds,
                               lastQuarkIndices,
                               hadronFlavour,
                               checkedHadronIds,
                               theLastQuarkIndex + 1)) {
        ambiguityResolved = true;
        break;
      };
    }
 
  }  // End of loop over all previous hadrons
 
  checkedHadronIds.erase(hadId);  // Removing the hadron from the list of checked hadrons
  if (nSameFlavourHadrons > 0 && !ambiguityResolved) {
    hadronFlavour.at(hadId) = qmFlav * 21;
    return true;
  }
  lastQuarkIndices.at(hadId) = lastQuarkIndex;
  hadronFlavour.at(hadId) = hadFlavour;
  return true;
}
 
//define this as a plug-in
DEFINE_FWK_MODULE(GenHFHadronMatcher);