RestOfEvent.h

/**************************************************************************
 * basf2 (Belle II Analysis Software Framework)                           *
 * Author: The Belle II Collaboration                                     *
 *                                                                        *
 * See git log for contributors and copyright holders.                    *
 * This file is licensed under LGPL-3.0, see LICENSE.md.                  *
 **************************************************************************/
 
#pragma once
#include <framework/datastore/RelationsObject.h>
#include <mdst/dataobjects/PIDLikelihood.h>
#include <analysis/VariableManager/Utility.h>
#include <analysis/dataobjects/Particle.h>
 
#include <framework/logging/Logger.h>
 
#include <Math/Vector4D.h>
 
#include <vector>
#include <string>
#include <set>
 
namespace Belle2 {
 
  // forward declarations
  class Particle;
  class ECLCluster;
  class KLMCluster;
  class Track;
 
  // TODO: Add support for the MdstVee dataobjects when they become available.
 

 
  class RestOfEvent : public RelationsObject {
 
  public:
    static constexpr const char* c_defaultMaskName = "all"; 
    struct Mask {
    public:
      Mask(const std::string& name = c_defaultMaskName, const std::string& origin = "unknown"): m_name(name),
        m_origin(origin)
      {
        B2DEBUG(10, "Mask " << name << " is being initialized by " << origin);
        m_isValid = false;
      };

      std::string getName() const
      {
        return m_name;
      }

      bool isValid() const
      {
        return m_isValid;
      }

      void addParticles(const std::vector<const Particle*>& particles)
      {
        if (isValid()) {
          B2INFO("Mask " + m_name + " originating from "  + m_origin + " is  valid, cannot write to it!");
          return;
        } else {
          for (auto* particle : particles) {
            m_maskedParticleIndices.insert(particle->getArrayIndex());
          }
          m_isValid = true;
        }
      }

      std::set<int> getParticles() const
      {
        return m_maskedParticleIndices;
      }

      std::set<int> getV0s() const
      {
        return m_maskedV0Indices;
      }

      void addV0(const Particle* v0, std::vector<int>& toErase)
      {
        m_maskedV0Indices.insert(v0->getArrayIndex());
        for (int& i : toErase) {
          m_maskedParticleIndices.erase(i);
        }
        m_maskedParticleIndices.insert(v0->getArrayIndex());
      }

      bool hasV0(const Particle* v0) const
      {
        return m_maskedV0Indices.count(v0->getArrayIndex()) > 0;
      }

 
      void clearParticles()
      {
        m_maskedParticleIndices.clear();
        m_maskedV0Indices.clear();
        m_isValid = false;
      }

      void print() const
      {
        B2INFO("Mask name: " + m_name + " originating from " + m_origin);
        if (!m_isValid) {
          B2INFO("\tNot valid!");
        }
        std::string printout =  "\tIndices: ";
        for (const int index : m_maskedParticleIndices) {
          printout += std::to_string(index) +  ", ";
        }
        B2INFO(printout);
      }
    private:
      std::string m_name;                       
      std::string m_origin;                     
      bool m_isValid;                           
      std::set<int> m_maskedParticleIndices;    
      std::set<int> m_maskedV0Indices;          
    };

    explicit RestOfEvent(int pdgCode = 0,
                         bool isNested = false,
                         bool isFromMC = false,
                         bool useKLMEnergy = false,
                         bool builtWithMostLikely = false):
      m_pdgCode(pdgCode), m_isNested(isNested), m_isFromMC(isFromMC), m_useKLMEnergy(useKLMEnergy),
      m_builtWithMostLikely(builtWithMostLikely) { };
    // setters
    void addParticles(const std::vector<const Particle*>& particle);
    void setPDGCode(int pdgCode)
    {
      m_pdgCode = pdgCode;
    }

    Particle* convertToParticle(const std::string& maskName = c_defaultMaskName, int pdgCode = 0, bool isSelfConjugated = true);
    int getPDGCode() const
    {
      return m_pdgCode;
    }

    bool hasParticle(const Particle* particle, const std::string& maskName = c_defaultMaskName) const;
    void initializeMask(const std::string& name, const std::string& origin = "unknown");
    void updateMaskWithCuts(const std::string& name, const std::shared_ptr<Variable::Cut>& trackCut = nullptr,
                            const std::shared_ptr<Variable::Cut>& eclCut = nullptr, const std::shared_ptr<Variable::Cut>& klmCut = nullptr,
                            bool updateExisting = false);
    void excludeParticlesFromMask(const std::string& maskName, const std::vector<const Particle*>& particles,
                                  Particle::EParticleSourceObject listType,
                                  bool discard);
    bool hasMask(const std::string& name) const;
    void updateMaskWithV0(const std::string& name, const Particle* particleV0);
    bool checkCompatibilityOfMaskAndV0(const std::string& name, const Particle* particleV0);
    bool getIsNested() const {return m_isNested;}

    bool isBuiltWithMostLikely() const {return m_builtWithMostLikely;}
    // getters
    std::vector<const Particle*> getParticles(const std::string& maskName = c_defaultMaskName, bool unpackComposite = true) const;
    std::vector<const Particle*> getPhotons(const std::string& maskName = c_defaultMaskName, bool unpackComposite = true) const;
    std::vector<const Particle*> getHadrons(const std::string& maskName = c_defaultMaskName, bool unpackComposite = true) const;
    std::vector<const Particle*> getChargedParticles(const std::string& maskName = c_defaultMaskName, unsigned int pdg = 0,
                                                     bool unpackComposite = true) const;

    ROOT::Math::PxPyPzEVector get4Vector(const std::string& maskName = c_defaultMaskName) const;

    ROOT::Math::PxPyPzEVector get4VectorNeutralECLClusters(const std::string& maskName = c_defaultMaskName) const;

    int getNTracks(const std::string& maskName = c_defaultMaskName) const;

    int getNECLClusters(const std::string& maskName = c_defaultMaskName) const;

    int getNKLMClusters(const std::string& maskName = c_defaultMaskName) const;

    std::vector<std::string> getMaskNames() const;

    void print(const std::string& maskName = c_defaultMaskName, bool unpackComposite = true) const;
 
  private:
 
    // persistent data members
    std::set<int> m_particleIndices;   
    std::vector<Mask> m_masks;         
    int m_pdgCode;                     
    bool m_isNested;                   
    bool m_isFromMC;                   
    bool m_useKLMEnergy;               
    bool m_builtWithMostLikely;        
 
    // Private methods
    bool isInParticleList(const Particle* roeParticle, const std::vector<const Particle*>& particlesToUpdate) const;

    Mask* findMask(const std::string& name);
    void printIndices(const std::string& maskName = c_defaultMaskName, bool unpackComposite = true,
                      const std::string& tab = " - ") const;
 
    ClassDef(RestOfEvent, 8) 
    // v8: added default mask name
    // v7: added m_builtWithMostLikely
 
  };
 

} // end namespace Belle2