NewV0Fitter.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/StoreArray.h>
#include <framework/gearbox/Const.h>
#include <mdst/dataobjects/Track.h>
#include <mdst/dataobjects/TrackFitResult.h>
#include <mdst/dataobjects/V0.h>
#include <tracking/dataobjects/V0ValidationVertex.h>
#include <tracking/dataobjects/RecoTrack.h>
#include <tracking/v0Finding/fitter/V0VertexFitter.h>
#include <genfit/Track.h>
#include <genfit/GFRaveVertex.h>
#include <map>
#include <memory>
#include <utility>
 
namespace Belle2 {

  class NewV0Fitter {
 
  public:

    struct FittedTrack {
      const RecoTrack* recoTrack = nullptr; 
      Const::ParticleType ptype = Const::invalidParticle; 
      genfit::MeasuredStateOnPlane state; 
      double pValue = 0; 
      int Ndf = 0; 

      void set(const RecoTrack* recoTrk, const Const::ParticleType& hypo, const genfit::MeasuredStateOnPlane& mSoP,
               const genfit::AbsTrackRep* rep)
      {
        recoTrack = recoTrk;
        ptype = hypo;
        state = mSoP;
        const auto* fitStatus = recoTrk->getGenfitTrack().getFitStatus(rep);
        pValue = fitStatus->getPVal();
        Ndf = fitStatus->getNdf();
      }
    };
 

    explicit NewV0Fitter(const std::string& trackFitResultsName = "",
                         const std::string& v0sName = "",
                         const std::string& v0ValidationVerticesName = "",
                         const std::string& recoTracksName = "",
                         const std::string& copiedRecoTracksName = "CopiedRecoTracks",
                         bool enableValidation = false);

    void initializeCuts(double vertexDistanceCut,
                        double vertexChi2Cut,
                        const std::tuple<double, double>& invMassRangeKshort,
                        const std::tuple<double, double>& invMassRangeLambda,
                        const std::tuple<double, double>& invMassRangePhoton);

    void setFitterMode(int fitterMode) {m_fitterMode = fitterMode;}

    void setVertexFitter(std::unique_ptr<V0VertexFitter> vertexFitter) {m_vertexFitter = std::move(vertexFitter);}

    bool fitAndStore(const Track* trackPlus, const Track* trackMinus, const Const::ParticleType& v0Hypothesis,
                     bool& isForceStored, bool& isHitRemoved);

    static std::pair<Const::ParticleType, Const::ParticleType> getTrackHypotheses(const Const::ParticleType& v0Hypothesis);
 
 
  private:

    enum ErrorStatus {
      c_NoTrackRepresentation = -4, 
      c_VertexFitFailed = -3,       
      c_ExtrapolationFailed = -2,   
      c_NotSelected = -1            
    };
 

    enum EInnerHitBits {
      c_BitTrackPlus = 0x1, 
      c_BitTrackMinus = 0x2 
    };
 

    int vertexFit(const RecoTrack* recoTrackPlus, const RecoTrack* recoTrackMinus,
                  int pdgTrackPlus, int pdgTrackMinus, const Const::ParticleType& v0Hypothesis);

    static const genfit::AbsTrackRep* getTrackRepresentation(const RecoTrack* recoTrack, int pdgCode);

    static bool setCardinalRep(genfit::Track& gfTrack, int pdgCode);

    static int extrapolateToVertex(genfit::MeasuredStateOnPlane& statePlus, genfit::MeasuredStateOnPlane& stateMinus,
                                   const genfit::GFRaveVertex& vertex);

    RecoTrack* copyRecoTrack(const RecoTrack* origRecoTrack, const genfit::MeasuredStateOnPlane& state);

    const RecoTrack* removeHitsAndRefit(const RecoTrack* origRecoTrack, const RecoTrack* lastRecoTrack,
                                        const Const::ParticleType& ptype);

    int isInnermostClusterShared(const RecoTrack* recoTrackPlus, const RecoTrack* recoTrackMinus);

    const TrackFitResult* saveTrackFitResult(const FittedTrack& trk, int sharedInnermostCluster);
 
 
    std::string m_recoTracksName;                    
    StoreArray<RecoTrack> m_recoTracks;              
    StoreArray<TrackFitResult> m_trackFitResults;    
    StoreArray<V0> m_v0s;                            
    StoreArray<V0ValidationVertex> m_validationV0s;  
    StoreArray<RecoTrack> m_copiedRecoTracks;        
 
    double m_vertexDistanceCut = 0;  
    double m_vertexChi2Cut = 0;   
    std::map<int, std::pair<double, double> > m_invMassCuts; 
 
    std::unique_ptr<V0VertexFitter> m_vertexFitter; 
    int m_fitterMode = 1;  
    bool m_validation = false; 
 
    // temporary storage
 
    genfit::GFRaveVertex m_fittedVertex; 
    double m_momentum = 0; 
    double m_invMass = 0; 
    FittedTrack m_trkPlus; 
    FittedTrack m_trkMinus; 
 
  };

}