#ifndef RecoAlgos_ConstrainedFitCandProducer_h
#define RecoAlgos_ConstrainedFitCandProducer_h
/* \class ConstrainedFitCandProducer
 *
 * \author Luca Lista, INFN
 *
 */
#include "FWCore/Framework/interface/stream/EDProducer.h"
#include "FWCore/Utilities/interface/InputTag.h"
#include "CommonTools/UtilAlgos/interface/ParameterAdapter.h"
#include "CommonTools/UtilAlgos/interface/EventSetupInitTrait.h"
#include "DataFormats/Candidate/interface/CandidateFwd.h"
#include "SimGeneral/HepPDTRecord/interface/PdtEntry.h"
#include <vector>

template <typename Fitter,
          typename InputCollection = reco::CandidateCollection,
          typename OutputCollection = InputCollection,
          typename Init = typename ::reco::modules::EventSetupInit<Fitter>::type>
class ConstrainedFitCandProducer : public edm::stream::EDProducer<> {
public:
  explicit ConstrainedFitCandProducer(const edm::ParameterSet&);

private:
  edm::EDGetTokenT<InputCollection> srcToken_;
  bool setLongLived_;
  bool setMassConstraint_;
  bool setPdgId_;
  int pdgId_;
  Init fitterInit_;
  Fitter fitter_;
  void produce(edm::Event&, const edm::EventSetup&) override;
};

#include "FWCore/ParameterSet/interface/ParameterSet.h"
#include "MagneticField/Records/interface/IdealMagneticFieldRecord.h"
#include "DataFormats/Common/interface/Handle.h"
#include "FWCore/Framework/interface/ESHandle.h"
#include "FWCore/Framework/interface/Event.h"
#include "DataFormats/Candidate/interface/VertexCompositeCandidate.h"
#include "FWCore/Framework/interface/EventSetup.h"
#include <algorithm>

template <typename Fitter, typename InputCollection, typename OutputCollection, typename Init>
ConstrainedFitCandProducer<Fitter, InputCollection, OutputCollection, Init>::ConstrainedFitCandProducer(
    const edm::ParameterSet& cfg)
    : srcToken_(consumes<InputCollection>(cfg.template getParameter<edm::InputTag>("src"))),
      setLongLived_(false),
      setMassConstraint_(false),
      setPdgId_(false),
      fitterInit_(consumesCollector()),
      fitter_(reco::modules::make<Fitter>(cfg)) {
  produces<OutputCollection>();
  std::string alias(cfg.getParameter<std::string>("@module_label"));
  const std::string setLongLived("setLongLived");
  std::vector<std::string> vBoolParams = cfg.template getParameterNamesForType<bool>();
  bool found = find(vBoolParams.begin(), vBoolParams.end(), setLongLived) != vBoolParams.end();
  if (found)
    setLongLived_ = cfg.template getParameter<bool>("setLongLived");
  const std::string setMassConstraint("setMassConstraint");
  found = find(vBoolParams.begin(), vBoolParams.end(), setMassConstraint) != vBoolParams.end();
  if (found)
    setMassConstraint_ = cfg.template getParameter<bool>("setMassConstraint");
  const std::string setPdgId("setPdgId");
  std::vector<std::string> vIntParams = cfg.getParameterNamesForType<int>();
  found = find(vIntParams.begin(), vIntParams.end(), setPdgId) != vIntParams.end();
  if (found) {
    setPdgId_ = true;
    pdgId_ = cfg.getParameter<int>("setPdgId");
  }
}

namespace reco {
  namespace fitHelper {
    template <typename C>
    struct Adder {
      static void add(C* c, std::unique_ptr<reco::VertexCompositeCandidate> t) { c->push_back(*t); }
    };

    template <typename T>
    struct Adder<edm::OwnVector<T> > {
      static void add(edm::OwnVector<T>* c, std::unique_ptr<reco::VertexCompositeCandidate> t) {
        c->push_back(std::move(t));
      }
    };

    template <typename C>
    inline void add(C* c, std::unique_ptr<reco::VertexCompositeCandidate> t) {
      Adder<C>::add(c, std::move(t));
    }
  }  // namespace fitHelper
}  // namespace reco

template <typename Fitter, typename InputCollection, typename OutputCollection, typename Init>
void ConstrainedFitCandProducer<Fitter, InputCollection, OutputCollection, Init>::produce(edm::Event& evt,
                                                                                          const edm::EventSetup& es) {
  fitterInit_.init(fitter_, evt, es);
  edm::Handle<InputCollection> cands;
  evt.getByToken(srcToken_, cands);
  auto fitted = std::make_unique<OutputCollection>();
  fitted->reserve(cands->size());
  for (typename InputCollection::const_iterator c = cands->begin(); c != cands->end(); ++c) {
    auto clone = std::make_unique<reco::VertexCompositeCandidate>(*c);
    fitter_.set(*clone);
    if (setLongLived_)
      clone->setLongLived();
    if (setMassConstraint_)
      clone->setMassConstraint();
    if (setPdgId_)
      clone->setPdgId(pdgId_);
    reco::fitHelper::add(fitted.get(), std::move(clone));
  }
  evt.put(std::move(fitted));
}

#endif
