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 <algorithm>


// user include files
#include "FWCore/Framework/interface/Frameworkfwd.h"
#include "FWCore/Framework/interface/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 "DataFormats/HepMCCandidate/interface/GenParticle.h"

#include "DataFormats/JetReco/interface/GenJet.h"
#include "DataFormats/Math/interface/deltaR.h"

#include "SimDataFormats/JetMatching/interface/JetFlavourInfoMatching.h"

#include "SimGeneral/HepPDTRecord/interface/ParticleDataTable.h"



//
// class declaration
//

class GenHFHadronMatcher : public edm::EDProducer
{
public:
    explicit GenHFHadronMatcher ( const edm::ParameterSet& );
    ~GenHFHadronMatcher();

    static void fillDescriptions ( edm::ConfigurationDescriptions& descriptions );

private:
    virtual void beginJob() ;
    virtual void produce ( edm::Event&, const edm::EventSetup& );
    virtual void endJob() ;

    virtual void beginRun ( edm::Run&, edm::EventSetup const& );
    virtual void endRun ( edm::Run&, edm::EventSetup const& );
    virtual void beginLuminosityBlock ( edm::LuminosityBlock&, edm::EventSetup const& );
    virtual void endLuminosityBlock ( edm::LuminosityBlock&, edm::EventSetup const& );

    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 );
    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 );
    bool putMotherIndex ( std::vector<std::vector<int> > &hadMothersIndices, int partIndex, int mothIndex );
    bool isHadron ( const int flavour, const reco::Candidate* thisParticle );
    bool isHadronPdgId ( const int flavour, const int pdgId );
    bool hasHadronDaughter ( const int flavour, const reco::Candidate* thisParticle );
    int isInList ( std::vector<const reco::Candidate*> particleList, const reco::Candidate* particle );
    int isInList ( std::vector<int> list, const int value );
    int findInMothers ( int idx, std::vector<int> &mothChains, std::vector<std::vector<int> > &hadMothersIndices, 
                        std::vector<reco::GenParticle> &hadMothers, int status, int pdgId, bool pdgAbs, int stopId, int firstLast, bool verbose );
    bool isNeutralPdg ( int pdgId );

    bool checkForLoop ( std::vector<const reco::Candidate*> &particleChain, const reco::Candidate* particle );
    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);

// ----------member data ---------------------------
    edm::EDGetTokenT<reco::GenParticleCollection> genParticlesToken_;
    edm::EDGetTokenT<reco::JetFlavourInfoMatchingCollection> jetFlavourInfosToken_;
    int flavour_;
    bool noBBbarResonances_;
    bool onlyJetClusteredHadrons_;

    std::string flavourStr_;  // Name of the flavour specified in config file

    edm::ESHandle<ParticleDataTable> pdt_;
};

//
// constants, enums and typedefs
//


//
// static data member definitions
//

//
// constructors and destructor
//
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_           = cfg.getParameter<int> ( "flavour" );
    noBBbarResonances_ = cfg.getParameter<bool> ( "noBBbarResonances" );
    onlyJetClusteredHadrons_ = cfg.getParameter<bool> ( "onlyJetClusteredHadrons" );
    
    flavour_ = abs ( flavour_ ); // Make flavour independent of sign given in configuration
    if ( flavour_==5 ) {
        flavourStr_="B";
    } else if ( flavour_==4 ) {
        flavourStr_="C";
    } else {
        edm::LogError ( "GenHFHadronMatcher" ) << "Flavour option must be 4 (c-jet) or 5 (b-jet), but is: " << flavour_ << ". Correct this!";
    }

    // Hadron matching products
    produces< std::vector<reco::GenParticle> > ( "gen"+flavourStr_+"HadPlusMothers" ); // All mothers in all decay chains above any hadron of specified flavour
    produces< std::vector< std::vector<int> > > ( "gen"+flavourStr_+"HadPlusMothersIndices" ); // Indices of mothers of each hadMother
    produces< std::vector<int> > ( "gen"+flavourStr_+"HadIndex" ); // Index of hadron in the vector of hadMothers
    produces< std::vector<int> > ( "gen"+flavourStr_+"HadFlavour" ); // PdgId of the first non-b(c) quark mother with sign corresponding to hadron charge
    produces< std::vector<int> > ( "gen"+flavourStr_+"HadJetIndex" ); // Index of genJet matched to each hadron by jet clustering algorithm
    produces< std::vector<int> > ( "gen"+flavourStr_+"HadLeptonIndex" ); // Index of lepton found among the hadron decay products in the list of mothers
    produces< std::vector<int> > ( "gen"+flavourStr_+"HadLeptonHadronIndex" ); // Index of hadron the lepton is associated to
    produces< std::vector<int> > ( "gen"+flavourStr_+"HadLeptonViaTau" ); // Whether lepton comes directly from hadron or via tau decay
    produces< std::vector<int> > ( "gen"+flavourStr_+"HadFromTopWeakDecay" ); // Tells whether the hadron appears in the chain after top decay
    produces< std::vector<int> > ( "gen"+flavourStr_+"HadBHadronId" ); // Index of a b-hadron which the current hadron comes from (for c-hadrons)

}

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::Event& evt, const edm::EventSetup& setup )
{

    setup.getData ( pdt_ );

    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::auto_ptr<std::vector<reco::GenParticle> > hadMothers ( new std::vector<reco::GenParticle> );
    std::auto_ptr<std::vector<std::vector<int> > > hadMothersIndices ( new std::vector<std::vector<int> > );
    std::auto_ptr<std::vector<int> > hadIndex ( new std::vector<int> );
    std::auto_ptr<std::vector<int> > hadFlavour ( new std::vector<int> );
    std::auto_ptr<std::vector<int> > hadJetIndex ( new std::vector<int> );
    std::auto_ptr<std::vector<int> > hadLeptonIndex ( new std::vector<int> );
    std::auto_ptr<std::vector<int> > hadLeptonHadIndex ( new std::vector<int> );
    std::auto_ptr<std::vector<int> > hadLeptonViaTau( new std::vector<int> );
    std::auto_ptr<std::vector<int> > hadFromTopWeakDecay ( new std::vector<int> );
    std::auto_ptr<std::vector<int> > hadBHadronId ( new std::vector<int> );
    
    *hadJetIndex = findHadronJets (genParticles.product(), jetFlavourInfos.product(), *hadIndex, *hadMothers, *hadMothersIndices, *hadLeptonIndex, *hadLeptonHadIndex, *hadLeptonViaTau, *hadFlavour, *hadFromTopWeakDecay, *hadBHadronId );

    // Putting products to the event
    evt.put ( hadMothers,         "gen"+flavourStr_+"HadPlusMothers" );
    evt.put ( hadMothersIndices,  "gen"+flavourStr_+"HadPlusMothersIndices" );
    evt.put ( hadIndex,           "gen"+flavourStr_+"HadIndex" );
    evt.put ( hadFlavour,         "gen"+flavourStr_+"HadFlavour" );
    evt.put ( hadJetIndex,        "gen"+flavourStr_+"HadJetIndex" );
    evt.put ( hadLeptonIndex,     "gen"+flavourStr_+"HadLeptonIndex" );
    evt.put ( hadLeptonHadIndex,  "gen"+flavourStr_+"HadLeptonHadronIndex" );
    evt.put ( hadLeptonViaTau,    "gen"+flavourStr_+"HadLeptonViaTau" );
    evt.put ( hadFromTopWeakDecay,"gen"+flavourStr_+"HadFromTopWeakDecay" );
    evt.put ( hadBHadronId,     "gen"+flavourStr_+"HadBHadronId" );
}

// ------------ method called once each job just before starting event loop  ------------
void GenHFHadronMatcher::beginJob()
{
}

// ------------ method called once each job just after ending the event loop  ------------
void GenHFHadronMatcher::endJob()
{
}

// ------------ method called when starting to processes a run  ------------
void GenHFHadronMatcher::beginRun ( edm::Run&, edm::EventSetup const& )
{
}

// ------------ method called when ending the processing of a run  ------------
void
GenHFHadronMatcher::endRun ( edm::Run&, edm::EventSetup const& )
{
}

// ------------ method called when starting to processes a luminosity block  ------------
void GenHFHadronMatcher::beginLuminosityBlock ( edm::LuminosityBlock&, edm::EventSetup const& )
{
}

// ------------ method called when ending the processing of a luminosity block  ------------
void GenHFHadronMatcher::endLuminosityBlock ( edm::LuminosityBlock&, edm::EventSetup const& )
{
}



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 )
{
    std::vector<int> hadJetIndex;
    std::vector<const reco::Candidate*> hadMothersCand;

    int topDaughterQId = -1;
    int topBarDaughterQId= -1;
    
    // Looping over all jets to get hadrons belonging to it
    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, 0, -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
        }
    // FIXME: Should be moved to a separate loop over jets after analysing all hadrons (in case leptons come from non-clustered hadrons)
        // Looping over all leptons associated with the jet
        const reco::GenParticleRefVector& leptonsInJet = jetInfo.getLeptons();
        for(reco::GenParticleRefVector::const_iterator lepton = leptonsInJet.begin(); lepton != leptonsInJet.end(); ++lepton){
            bool leptonViaTau = false;
            // Skipping tau leptons
            if(std::abs((*lepton)->pdgId()) == 15) continue;
            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(!isHadron(flavour_, leptonMother)) continue;
            // Finding the index of this hadron in the list of analysed particles
            int leptonHadronParticleIndex = std::find(hadMothersCand.begin(), hadMothersCand.end(), leptonMother) - hadMothersCand.begin();
            if(leptonHadronParticleIndex >= (int)hadMothersCand.size()) continue;
            if(leptonHadronParticleIndex < 0) continue;
            // Finding the actual hadron index among those that were found
            int leptonHadronIndex = std::find(hadIndex.begin(), hadIndex.end(), leptonHadronParticleIndex) - hadIndex.begin();
            if(leptonHadronIndex >= (int)hadIndex.size()) continue;
            // Putting the lepton, its index and hadron index to the corresponding lists
            hadMothersCand.push_back(&**lepton);
            const int leptonIndex = hadMothersCand.size()-1;
            hadLeptonIndex.push_back(leptonIndex);
            hadLeptonViaTau.push_back((int)leptonViaTau);
            hadLeptonHadIndex.push_back(leptonHadronIndex);
        } 
    }       // 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, 0, -1 );
            // Storing the index of the hadron to the list
            hadIndex.push_back ( hadronIndex );
            hadJetIndex.push_back ( -1 );  // Jet index undefined
        }
    }


    for ( int i=0; i< ( int ) hadMothersCand.size(); i++ ) {
        hadMothers.push_back ( ( *dynamic_cast<const reco::GenParticle*> ( hadMothersCand.at(i) ) ) );
    }

    // 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;

        int hadFlav = hadMothers.at(hadIdx).pdgId() <0?-1:1; // Charge of the hadron (-1,1)
        if ( abs ( hadMothers.at(hadIdx).pdgId() ) /1000 < 1 ) {
            hadFlav*=-1;    // Inverting flavour of hadron if it is a meson
        }

        // Searching only first quark in the chain with the same flavour as hadron
        findInMothers ( hadIdx, FirstQuarkId, hadMothersIndices, hadMothers, 0, hadFlav*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
            int bQFlav = hadMothers.at(FirstQuarkId.at(qId)).pdgId() < 0?-1:1;
            // Finding last quark of the hadron starting from the first quark
            findInMothers ( FirstQuarkId.at(qId), LastQuarkId, hadMothersIndices, hadMothers, 0, bQFlav*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.size()<1) {
            hadronFlavour = 0;
        } else {
            int qIdx = LastQuarkId.at( lastQuarkIndices.at(hadNum) );
            int qFlav = ( hadMothers.at(qIdx).pdgId() < 0 ) ? -1 : 1;
            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.size()>0) {
            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::isInList ( std::vector<const reco::Candidate*> particleList, const reco::Candidate* particle )
{
    for ( unsigned int i = 0; i<particleList.size(); i++ )
        if ( particleList.at(i)==particle ) {
            return i;
        }

    return -1;
}

int GenHFHadronMatcher::isInList ( std::vector<int> list, const int value )
{
    for ( unsigned int i = 0; i<list.size(); i++ )
        if ( list.at(i)==value ) {
            return i;
        }

    return -1;
}


bool GenHFHadronMatcher::isHadron ( const int flavour, const reco::Candidate* thisParticle )
{
    return isHadronPdgId(flavour, thisParticle->pdgId());
}


bool GenHFHadronMatcher::isHadronPdgId ( const int flavour, const int pdgId )
{
    int flavour_abs = std::abs(flavour);
    if(flavour_abs > 5 || flavour_abs < 1) return false;
    int pdgId_abs = std::abs(pdgId);

    if ( pdgId_abs / 1000 == flavour_abs // baryons
            || ( pdgId_abs / 100 % 10 == flavour_abs // mesons
                 && ! ( noBBbarResonances_ && pdgId_abs / 10 % 100 == 11*flavour_abs ) // but not a resonance
               )
       ) {
        return true;
    } else {
        return false;
    }
}


bool GenHFHadronMatcher::hasHadronDaughter ( const int flavour, const reco::Candidate* thisParticle )
{
// 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 )
{

    // 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 = isInList ( 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 = isInList ( 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 = isInList ( 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 )
{
    // 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, std::vector<std::vector<int> > &hadMothersIndices, std::vector<reco::GenParticle> &hadMothers, int status, int pdgId, bool pdgAbs=false, int stopId=-1, int firstLast=0, bool verbose=false)
{
    int foundStopId = -1;
    int pdg_1 = hadMothers.at ( idx ).pdgId();
    int partCharge = ( hadMothers.at ( idx ).pdgId() >0 ) ?1:-1;
// Inverting charge if mother is a b(c) meson
    if ( abs ( hadMothers.at ( idx ).pdgId() ) /1000 < 1 && ( abs ( hadMothers.at ( idx ).pdgId() ) /100%10 == 4 || abs ( hadMothers.at ( idx ).pdgId() ) /100%10 == 5 ) ) {
        partCharge*=-1;
    }

    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 ( 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 )
            || ( abs ( hadMothers.at ( idx ).pdgId() ) == abs ( pdgId ) && pdgAbs==true ) )
            && ( hadMothers.at ( idx ).status() == status || status==0 )
            && hasCorrectMothers ) {
        isCorrect=true;
        bool inList=false;
        for ( unsigned int k=0; k<mothChains.size(); k++ ) if ( mothChains[k]==idx ) {
                inList=true;    // Checking whether isn't already in the list
                break;
            }
        if ( !inList && mothers.at ( 0 ) >=0 && ( hadMothers.at ( idx ).pdgId() *pdgId>0 || !pdgAbs ) ) {       // If not in list and mother of this quark has correct charge
            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[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 ) != 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 int max = 5;
    int neutralPdgs[max]= {9,21,22,23,25};
    for ( int i=0; i<max; i++ ) if ( abs ( pdgId ) ==neutralPdgs[i] ) {
            return true;
        }
    return false;
}


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)
{
    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 );