Putting everything together

#pragma once

#include "extensions/jana/JOmniFactory.h"

#include "algorithms/reco/ElectronReconstruction.h"


namespace eicrecon {

class ReconstructedElectrons_factory : public JOmniFactory<ReconstructedElectrons_factory, ElectronReconstructionConfig> {
private:

    // Underlying algorithm
    std::unique_ptr<eicrecon::ElectronReconstruction> m_algo;

    // Declare inputs
    PodioInput<edm4hep::MCParticle> m_in_mc_particles {this, "MCParticles"};
    PodioInput<edm4eic::ReconstructedParticle> m_in_rc_particles {this, "ReconstructedChargedParticles"};
    PodioInput<edm4eic::MCRecoParticleAssociation> m_in_rc_particles_assoc {this, "ReconstructedChargedParticleAssociations"};

    VariadicPodioInput<edm4eic::MCRecoClusterParticleAssociation> m_in_clu_assoc {this};

    // Declare outputs
    PodioOutput<edm4eic::ReconstructedParticle> m_out_reco_particles {this};

    // Declare parameters
    ParameterRef<double> m_min_energy_over_momentum {this, "minEnergyOverMomentum", config().min_energy_over_momentum};
    ParameterRef<double> m_max_energy_over_momentum {this, "maxEnergyOverMomentum", config().max_energy_over_momentum};

    // Declare services here, e.g.
    // Service<DD4hep_service> m_geoSvc {this};

public:
    void Configure() {
        // This is called when the factory is instantiated.
        // Use this callback to make sure the algorithm is configured.
        // The logger, parameters, and services have all been fetched before this is called
        m_algo = std::make_unique<eicrecon::ElectronReconstruction>();

        // Pass config object to algorithm
        m_algo->applyConfig(config());

        // If we needed geometry, we'd obtain it like so
        // m_algo->init(m_geoSvc().detector(), m_geoSvc().converter(), logger());

        m_algo->init(logger());
    }

    void ChangeRun(int64_t run_number) {
        // This is called whenever the run number is changed.
        // Use this callback to retrieve state that is keyed off of run number.
        // This state should usually be managed by a Service.
        // Note: You usually don't need this, because you can declare a Resource instead.
    }

    void Process(int64_t run_number, uint64_t event_number) {
        // This is called on every event.
        // Use this callback to call your Algorithm using all inputs and outputs
        // The inputs will have already been fetched for you at this point.
        auto output = m_algo->execute(
          m_in_mc_particles(),
          m_in_rc_particles(),
          m_in_rc_particles_assoc(),
          m_in_clu_assoc()
        );

        logger()->debug( "Event {}: Found {} reconstructed electron candidates", event_number, output->size() );

        m_out_reco_particles() = std::move(output);
        // JANA will take care of publishing the outputs for you.
    }
};
} // namespace eicrecon

    app->Add(new JOmniFactoryGeneratorT<ReconstructedElectrons_factory>(
        "ReconstructedElectrons",
        {"MCParticles", "ReconstructedChargedParticles", "ReconstructedChargedParticleAssociations",
        "EcalBarrelScFiClusterAssociations",
        "EcalEndcapNClusterAssociations",
        "EcalEndcapPClusterAssociations",
        "EcalEndcapPInsertClusterAssociations",
        "EcalLumiSpecClusterAssociations",
        },
        {"ReconstructedElectrons"},
        {},  // Override config values here
        app
    ));

    "ReconstructedElectrons",

#pragma once

#include <edm4eic/MCRecoClusterParticleAssociationCollection.h>
#include <edm4eic/MCRecoParticleAssociationCollection.h>
#include <edm4eic/ReconstructedParticleCollection.h>
#include <edm4hep/MCParticleCollection.h>
#include <spdlog/logger.h>
#include <memory>
#include <vector>

#include "ElectronReconstructionConfig.h"
#include "algorithms/interfaces/WithPodConfig.h"

namespace eicrecon {

    class ElectronReconstruction : public WithPodConfig<ElectronReconstructionConfig>{
    public:

        // Initialization will set the pointer of the logger
        void init(std::shared_ptr<spdlog::logger> logger);

        // Algorithm will apply E/p cut to reconstructed tracks based on truth track-cluster associations
        std::unique_ptr<edm4eic::ReconstructedParticleCollection> execute(
                const edm4hep::MCParticleCollection *mcparts,
                const edm4eic::ReconstructedParticleCollection *rcparts,
                const edm4eic::MCRecoParticleAssociationCollection *rcassoc,
                const std::vector<const edm4eic::MCRecoClusterParticleAssociationCollection*> &in_clu_assoc
        );
        
        // Could overload execute here to, for instance, use track projections 
        // for track-cluster association (instead of truth information)
        
    private:
        std::shared_ptr<spdlog::logger> m_log;  // pointer to logger
        double m_electron{0.000510998928};      // electron mass
    };
} // namespace eicrecon

struct ElectronReconstructionConfig {

    double min_energy_over_momentum = 0.9;
    double max_energy_over_momentum = 1.2;

};

#include "ElectronReconstruction.h"

#include <edm4eic/ClusterCollection.h>
#include <edm4hep/utils/vector_utils.h>
#include <fmt/core.h>

namespace eicrecon {

  void ElectronReconstruction::init(std::shared_ptr<spdlog::logger> logger) {
    m_log = logger;
  }

  std::unique_ptr<edm4eic::ReconstructedParticleCollection> ElectronReconstruction::execute(
    const edm4hep::MCParticleCollection *mcparts,
    const edm4eic::ReconstructedParticleCollection *rcparts,
    const edm4eic::MCRecoParticleAssociationCollection *rcassoc,
    const std::vector<const edm4eic::MCRecoClusterParticleAssociationCollection*> &in_clu_assoc
    ) {

        // Step 1. Loop through MCParticle - cluster associations
        // Step 2. Get Reco particle for the Mc Particle matched to cluster
        // Step 3. Apply E/p cut using Reco cluster Energy and Reco Particle momentum

        // Some obvious improvements:
        // - E/p cut from real study optimized for electron finding and hadron rejection
        // - use of any HCAL info?
        // - check for duplicates?

        // output container
        auto out_electrons = std::make_unique<edm4eic::ReconstructedParticleCollection>();

        for ( const auto *col : in_clu_assoc ){ // loop on cluster association collections
          for ( auto clu_assoc : (*col) ){ // loop on MCRecoClusterParticleAssociation in this particular collection
            auto sim = clu_assoc.getSim(); // McParticle
            auto clu = clu_assoc.getRec(); // RecoCluster

            m_log->trace( "SimId={}, CluId={}", clu_assoc.getSimID(), clu_assoc.getRecID() );
            m_log->trace( "MCParticle: Energy={} GeV, p={} GeV, E/p = {} for PDG: {}", clu.getEnergy(), edm4hep::utils::magnitude(sim.getMomentum()), clu.getEnergy() / edm4hep::utils::magnitude(sim.getMomentum()), sim.getPDG() );


            // Find the Reconstructed particle associated to the MC Particle that is matched with this reco cluster
            // i.e. take (MC Particle <-> RC Cluster) + ( MC Particle <-> RC Particle ) = ( RC Particle <-> RC Cluster )
            auto reco_part_assoc = rcassoc->begin();
            for (; reco_part_assoc != rcassoc->end(); ++reco_part_assoc) {
              if (reco_part_assoc->getSimID() == (unsigned) clu_assoc.getSimID()) {
                break;
              }
            }

            // if we found a reco particle then test for electron compatibility
            if ( reco_part_assoc != rcassoc->end() ){
              auto reco_part = reco_part_assoc->getRec();
              double EoverP = clu.getEnergy() / edm4hep::utils::magnitude(reco_part.getMomentum());
              m_log->trace( "ReconstructedParticle: Energy={} GeV, p={} GeV, E/p = {} for PDG (from truth): {}", clu.getEnergy(), edm4hep::utils::magnitude(reco_part.getMomentum()), EoverP, sim.getPDG() );

              // Apply the E/p cut here to select electons
              if ( EoverP >= m_cfg.min_energy_over_momentum && EoverP <= m_cfg.max_energy_over_momentum ) {
                out_electrons->push_back(reco_part.clone());
              }

            } else {
              m_log->debug( "Could not find reconstructed particle for SimId={}", clu_assoc.getSimID() );
            }

          } // loop on MC particle to cluster associations in collection
        } // loop on collections

        m_log->debug( "Found {} electron candidates", out_electrons->size() );
        return out_electrons;
    }

}