#ifndef CommonTools_RecoAlgos_PrimaryVertexSorter_
#define CommonTools_RecoAlgos_PrimaryVertexSorter_

// system include files
#include <memory>
#include <string>

// user include files
#include "FWCore/Framework/interface/Frameworkfwd.h"
#include "FWCore/Framework/interface/stream/EDProducer.h"
#include "FWCore/ParameterSet/interface/ParameterSet.h"

#include "DataFormats/Common/interface/ValueMap.h"
#include "DataFormats/Common/interface/Association.h"

#include "FWCore/Framework/interface/Event.h"
#include "FWCore/Framework/interface/MakerMacros.h"

#include "DataFormats/VertexReco/interface/VertexFwd.h"

#include "CommonTools/RecoAlgos/interface/PrimaryVertexAssignment.h"
#include "CommonTools/RecoAlgos/interface/PrimaryVertexSorting.h"
#include "DataFormats/PatCandidates/interface/PackedCandidate.h"

#include "TrackingTools/TransientTrack/interface/TransientTrackBuilder.h"
#include "TrackingTools/Records/interface/TransientTrackRecord.h"

/**\class PrimaryVertexSorter
 * \author Andrea Rizzi

*/

template <class ParticlesCollection>

class PrimaryVertexSorter : public edm::stream::EDProducer<> {
public:
  typedef edm::Association<reco::VertexCollection> CandToVertex;
  typedef edm::ValueMap<int> CandToVertexQuality;
  typedef edm::ValueMap<float> VertexScore;

  typedef ParticlesCollection PFCollection;

  explicit PrimaryVertexSorter(const edm::ParameterSet&);

  ~PrimaryVertexSorter() override {}

  void produce(edm::Event&, const edm::EventSetup&) override;

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

private:
  PrimaryVertexAssignment assignmentAlgo_;
  PrimaryVertexSorting sortingAlgo_;

  /// Candidates to be analyzed
  edm::EDGetTokenT<PFCollection> tokenCandidates_;

  /// vertices
  edm::EDGetTokenT<reco::VertexCollection> tokenVertices_;
  edm::EDGetTokenT<edm::View<reco::Candidate>> tokenJets_;
  edm::EDGetTokenT<edm::ValueMap<float>> tokenTrackTimeTag_;
  edm::EDGetTokenT<edm::ValueMap<float>> tokenTrackTimeResoTag_;

  edm::ESGetToken<TransientTrackBuilder, TransientTrackRecord> tokenBuilder_;

  bool produceOriginalMapping_;
  bool produceSortedVertices_;
  bool producePFPileUp_;
  bool producePFNoPileUp_;
  int qualityCut_;
  bool useMET_;
  bool useTiming_;

  static void fillDescriptionsCommon(edm::ParameterSetDescription& descriptions);

  void doConsumesForTiming(const edm::ParameterSet& iConfig);
  bool needsProductsForTiming();
  std::pair<int, PrimaryVertexAssignment::Quality> runAlgo(const reco::VertexCollection& vertices,
                                                           const typename ParticlesCollection::value_type& pf,
                                                           const edm::ValueMap<float>* trackTimeTag,
                                                           const edm::ValueMap<float>* trackTimeResoTag,
                                                           const edm::View<reco::Candidate>& jets,
                                                           const TransientTrackBuilder& builder);
};

#include "DataFormats/VertexReco/interface/Vertex.h"
#include "FWCore/Framework/interface/ESHandle.h"

// #include "FWCore/MessageLogger/interface/MessageLogger.h"
#include "FWCore/Utilities/interface/Exception.h"
#include "FWCore/Framework/interface/EventSetup.h"

template <class ParticlesCollection>
PrimaryVertexSorter<ParticlesCollection>::PrimaryVertexSorter(const edm::ParameterSet& iConfig)
    : assignmentAlgo_(iConfig.getParameterSet("assignment")),
      sortingAlgo_(iConfig.getParameterSet("sorting")),
      tokenCandidates_(consumes<ParticlesCollection>(iConfig.getParameter<edm::InputTag>("particles"))),
      tokenVertices_(consumes<reco::VertexCollection>(iConfig.getParameter<edm::InputTag>("vertices"))),
      tokenJets_(consumes<edm::View<reco::Candidate>>(iConfig.getParameter<edm::InputTag>("jets"))),
      tokenBuilder_(esConsumes(edm::ESInputTag("", "TransientTrackBuilder"))),
      produceOriginalMapping_(iConfig.getParameter<bool>("produceAssociationToOriginalVertices")),
      produceSortedVertices_(iConfig.getParameter<bool>("produceSortedVertices")),
      producePFPileUp_(iConfig.getParameter<bool>("producePileUpCollection")),
      producePFNoPileUp_(iConfig.getParameter<bool>("produceNoPileUpCollection")),
      qualityCut_(iConfig.getParameter<int>("qualityForPrimary")),
      useMET_(iConfig.getParameter<bool>("usePVMET")),
      useTiming_(iConfig.getParameterSet("assignment").getParameter<bool>("useTiming")) {
  using namespace std;
  using namespace edm;
  using namespace reco;

  if (produceOriginalMapping_) {
    produces<CandToVertex>("original");
    produces<CandToVertexQuality>("original");
    produces<VertexScore>("original");
  }
  if (produceSortedVertices_) {
    produces<reco::VertexCollection>();
    produces<CandToVertex>();
    produces<CandToVertexQuality>();
    produces<VertexScore>();
  }

  if (producePFPileUp_) {
    if (produceOriginalMapping_)
      produces<PFCollection>("originalPileUp");
    if (produceSortedVertices_)
      produces<PFCollection>("PileUp");
  }

  if (producePFNoPileUp_) {
    if (produceOriginalMapping_)
      produces<PFCollection>("originalNoPileUp");
    if (produceSortedVertices_)
      produces<PFCollection>("NoPileUp");
  }

  if (useTiming_)
    doConsumesForTiming(iConfig);
}

template <class ParticlesCollection>
void PrimaryVertexSorter<ParticlesCollection>::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
  using namespace std;
  using namespace edm;
  using namespace reco;

  Handle<edm::View<reco::Candidate>> jets;
  iEvent.getByToken(tokenJets_, jets);

  TransientTrackBuilder const& builder = iSetup.getData(tokenBuilder_);

  Handle<VertexCollection> vertices;
  iEvent.getByToken(tokenVertices_, vertices);

  Handle<ParticlesCollection> particlesHandle;
  iEvent.getByToken(tokenCandidates_, particlesHandle);

  Handle<edm::ValueMap<float>> trackTimeTagHandle;
  Handle<edm::ValueMap<float>> trackTimeResoTagHandle;

  const edm::ValueMap<float>* trackTimeTag = nullptr;
  const edm::ValueMap<float>* trackTimeResoTag = nullptr;
  if (useTiming_ && needsProductsForTiming()) {
    iEvent.getByToken(tokenTrackTimeTag_, trackTimeTagHandle);
    iEvent.getByToken(tokenTrackTimeResoTag_, trackTimeResoTagHandle);

    trackTimeTag = trackTimeTagHandle.product();
    trackTimeResoTag = trackTimeResoTagHandle.product();
  }

  ParticlesCollection particles = *particlesHandle.product();
  std::vector<int> pfToPVVector;
  std::vector<PrimaryVertexAssignment::Quality> pfToPVQualityVector;
  //reverse mapping
  std::vector<std::vector<int>> pvToPFVector(vertices->size());
  std::vector<std::vector<const reco::Candidate*>> pvToCandVector(vertices->size());
  std::vector<std::vector<PrimaryVertexAssignment::Quality>> pvToPFQualityVector(vertices->size());
  std::vector<float> vertexScoreOriginal(vertices->size());
  std::vector<float> vertexScore(vertices->size());

  for (auto const& pf : particles) {
    std::pair<int, PrimaryVertexAssignment::Quality> vtxWithQuality =
        runAlgo(*vertices, pf, trackTimeTag, trackTimeResoTag, *jets, builder);
    pfToPVVector.push_back(vtxWithQuality.first);
    pfToPVQualityVector.push_back(vtxWithQuality.second);
  }

  //Invert the mapping
  for (size_t i = 0; i < pfToPVVector.size(); i++) {
    auto pv = pfToPVVector[i];
    auto qual = pfToPVQualityVector[i];
    if (pv >= 0 and qual >= qualityCut_) {
      pvToPFVector[pv].push_back(i);
      //    std::cout << i << std::endl;
      //     const typename  ParticlesCollection::value_type & cp = particles[i];
      //     std::cout << "CP " << &cp <<  std::endl;
      pvToCandVector[pv].push_back(&particles[i]);
      pvToPFQualityVector[pv].push_back(qual);
    }
  }

  //Use multimap for sorting of indices
  std::multimap<float, int> scores;
  for (unsigned int i = 0; i < vertices->size(); i++) {
    float s = sortingAlgo_.score((*vertices)[i], pvToCandVector[i], useMET_);
    vertexScoreOriginal[i] = s;
    scores.insert(std::pair<float, int>(-s, i));
  }

  //create indices
  std::vector<int> oldToNew(vertices->size()), newToOld(vertices->size());
  size_t newIdx = 0;
  for (auto const& idx : scores) {
    //    std::cout << newIdx << " score: " << idx.first << " oldidx: " << idx.second << " "<< producePFPileUp_ << std::endl;
    vertexScore[newIdx] = -idx.first;
    oldToNew[idx.second] = newIdx;
    newToOld[newIdx] = idx.second;
    newIdx++;
  }

  if (produceOriginalMapping_) {
    unique_ptr<CandToVertex> pfCandToOriginalVertexOutput(new CandToVertex(vertices));
    unique_ptr<CandToVertexQuality> pfCandToOriginalVertexQualityOutput(new CandToVertexQuality());
    CandToVertex::Filler cand2VertexFiller(*pfCandToOriginalVertexOutput);
    CandToVertexQuality::Filler cand2VertexQualityFiller(*pfCandToOriginalVertexQualityOutput);

    cand2VertexFiller.insert(particlesHandle, pfToPVVector.begin(), pfToPVVector.end());
    cand2VertexQualityFiller.insert(particlesHandle, pfToPVQualityVector.begin(), pfToPVQualityVector.end());

    cand2VertexFiller.fill();
    cand2VertexQualityFiller.fill();
    iEvent.put(std::move(pfCandToOriginalVertexOutput), "original");
    iEvent.put(std::move(pfCandToOriginalVertexQualityOutput), "original");

    unique_ptr<VertexScore> vertexScoreOriginalOutput(new VertexScore);
    VertexScore::Filler vertexScoreOriginalFiller(*vertexScoreOriginalOutput);
    vertexScoreOriginalFiller.insert(vertices, vertexScoreOriginal.begin(), vertexScoreOriginal.end());
    vertexScoreOriginalFiller.fill();
    iEvent.put(std::move(vertexScoreOriginalOutput), "original");
  }

  if (produceSortedVertices_) {
    std::vector<int> pfToSortedPVVector;
    //      std::vector<int> pfToSortedPVQualityVector;
    for (size_t i = 0; i < pfToPVVector.size(); i++) {
      pfToSortedPVVector.push_back(oldToNew[pfToPVVector[i]]);
      //        pfToSortedPVQualityVector.push_back(pfToPVQualityVector[i]); //same as old!
    }

    unique_ptr<reco::VertexCollection> sortedVerticesOutput(new reco::VertexCollection);
    for (size_t i = 0; i < vertices->size(); i++) {
      sortedVerticesOutput->push_back((*vertices)[newToOld[i]]);
    }
    edm::OrphanHandle<reco::VertexCollection> oh = iEvent.put(std::move(sortedVerticesOutput));
    unique_ptr<CandToVertex> pfCandToVertexOutput(new CandToVertex(oh));
    unique_ptr<CandToVertexQuality> pfCandToVertexQualityOutput(new CandToVertexQuality());
    CandToVertex::Filler cand2VertexFiller(*pfCandToVertexOutput);
    CandToVertexQuality::Filler cand2VertexQualityFiller(*pfCandToVertexQualityOutput);

    cand2VertexFiller.insert(particlesHandle, pfToSortedPVVector.begin(), pfToSortedPVVector.end());
    cand2VertexQualityFiller.insert(particlesHandle, pfToPVQualityVector.begin(), pfToPVQualityVector.end());

    cand2VertexFiller.fill();
    cand2VertexQualityFiller.fill();
    iEvent.put(std::move(pfCandToVertexOutput));
    iEvent.put(std::move(pfCandToVertexQualityOutput));

    unique_ptr<VertexScore> vertexScoreOutput(new VertexScore);
    VertexScore::Filler vertexScoreFiller(*vertexScoreOutput);
    vertexScoreFiller.insert(oh, vertexScore.begin(), vertexScore.end());
    vertexScoreFiller.fill();
    iEvent.put(std::move(vertexScoreOutput));
  }

  unique_ptr<PFCollection> pfCollectionNOPUOriginalOutput(new PFCollection);
  unique_ptr<PFCollection> pfCollectionNOPUOutput(new PFCollection);
  unique_ptr<PFCollection> pfCollectionPUOriginalOutput(new PFCollection);
  unique_ptr<PFCollection> pfCollectionPUOutput(new PFCollection);

  for (size_t i = 0; i < particles.size(); i++) {
    auto pv = pfToPVVector[i];
    auto qual = pfToPVQualityVector[i];

    if (producePFNoPileUp_ && produceSortedVertices_)
      if (pv == newToOld[0] and qual >= qualityCut_)
        pfCollectionNOPUOutput->push_back(particles[i]);

    if (producePFPileUp_ && produceSortedVertices_)
      if (pv != newToOld[0] and qual >= qualityCut_)
        pfCollectionPUOutput->push_back(particles[i]);

    if (producePFNoPileUp_ && produceOriginalMapping_)
      if (pv == 0 and qual >= qualityCut_)
        pfCollectionNOPUOriginalOutput->push_back(particles[i]);

    if (producePFPileUp_ && produceOriginalMapping_)
      if (pv != 0 and qual >= qualityCut_)
        pfCollectionPUOriginalOutput->push_back(particles[i]);
  }
  if (producePFNoPileUp_ && produceSortedVertices_)
    iEvent.put(std::move(pfCollectionNOPUOutput), "NoPileUp");
  if (producePFPileUp_ && produceSortedVertices_)
    iEvent.put(std::move(pfCollectionPUOutput), "PileUp");
  if (producePFNoPileUp_ && produceOriginalMapping_)
    iEvent.put(std::move(pfCollectionNOPUOriginalOutput), "originalNoPileUp");
  if (producePFPileUp_ && produceOriginalMapping_)
    iEvent.put(std::move(pfCollectionPUOriginalOutput), "originalPileUp");
}

template <>
inline void PrimaryVertexSorter<std::vector<reco::RecoChargedRefCandidate>>::fillDescriptions(
    edm::ConfigurationDescriptions& descriptions) {
  edm::ParameterSetDescription desc;
  fillDescriptionsCommon(desc);
  desc.add<edm::InputTag>("trackTimeTag", edm::InputTag(""));
  desc.add<edm::InputTag>("trackTimeResoTag", edm::InputTag(""));
  desc.add<edm::InputTag>("particles", edm::InputTag("trackRefsForJets"));
  desc.add<edm::InputTag>("vertices", edm::InputTag("offlinePrimaryVertices"));
  desc.add<edm::InputTag>("jets", edm::InputTag("ak4CaloJetsForTrk"));
  desc.add<bool>("produceAssociationToOriginalVertices", false);
  desc.add<bool>("produceSortedVertices", true);
  desc.add<bool>("producePileUpCollection", false);
  desc.add<bool>("produceNoPileUpCollection", false);
  descriptions.add("sortedPrimaryVertices", desc);
}

template <>
inline void PrimaryVertexSorter<std::vector<reco::PFCandidate>>::fillDescriptions(
    edm::ConfigurationDescriptions& descriptions) {
  edm::ParameterSetDescription desc;
  fillDescriptionsCommon(desc);
  desc.add<edm::InputTag>("particles", edm::InputTag("particleFlow"));
  desc.add<edm::InputTag>("vertices", edm::InputTag("offlinePrimaryVertices"));
  desc.add<edm::InputTag>("jets", edm::InputTag("ak4PFJets"));
  desc.add<bool>("produceAssociationToOriginalVertices", true);
  desc.add<bool>("produceSortedVertices", true);
  desc.add<bool>("producePileUpCollection", true);
  desc.add<bool>("produceNoPileUpCollection", true);
  descriptions.add("sortedPFPrimaryVertices", desc);
}

template <>
inline void PrimaryVertexSorter<std::vector<pat::PackedCandidate>>::fillDescriptions(
    edm::ConfigurationDescriptions& descriptions) {
  edm::ParameterSetDescription desc;
  fillDescriptionsCommon(desc);
  desc.add<edm::InputTag>("particles", edm::InputTag("packedPFCandidates"));
  desc.add<edm::InputTag>("vertices", edm::InputTag("offlineSlimmedPrimaryVertices"));
  desc.add<edm::InputTag>("jets", edm::InputTag("slimmedJets"));
  desc.add<bool>("produceAssociationToOriginalVertices", true);
  desc.add<bool>("produceSortedVertices", true);
  desc.add<bool>("producePileUpCollection", true);
  desc.add<bool>("produceNoPileUpCollection", true);
  descriptions.add("sortedPackedPrimaryVertices", desc);
}

template <class ParticlesCollection>
inline void PrimaryVertexSorter<ParticlesCollection>::fillDescriptionsCommon(edm::ParameterSetDescription& desc) {
  {
    edm::ParameterSetDescription psd0;
    desc.add<edm::ParameterSetDescription>("sorting", psd0);
  }
  {
    edm::ParameterSetDescription psd0;
    psd0.add<double>("maxDzSigForPrimaryAssignment", 5.0);
    psd0.add<double>("maxDzForPrimaryAssignment", 0.1);
    psd0.add<double>("maxDzErrorForPrimaryAssignment", 0.05);
    psd0.add<double>("maxDtSigForPrimaryAssignment", 3.0);
    psd0.add<double>("maxJetDeltaR", 0.5);
    psd0.add<double>("minJetPt", 25);
    psd0.add<double>("maxDistanceToJetAxis", 0.07);
    psd0.add<double>("maxDzForJetAxisAssigment", 0.1);
    psd0.add<double>("maxDxyForJetAxisAssigment", 0.1);
    psd0.add<double>("maxDxySigForNotReconstructedPrimary", 2);
    psd0.add<double>("maxDxyForNotReconstructedPrimary", 0.01);
    psd0.add<bool>("useTiming", false);
    psd0.add<bool>("useVertexFit", true);
    psd0.add<bool>("preferHighRanked", false);
    psd0.add<unsigned int>("NumOfPUVtxsForCharged", 0);
    psd0.add<double>("DzCutForChargedFromPUVtxs", 0.2);
    psd0.add<double>("PtMaxCharged", -1);
    psd0.add<double>("EtaMinUseDz", -1);
    psd0.add<bool>("OnlyUseFirstDz", false);
    desc.add<edm::ParameterSetDescription>("assignment", psd0);
  }
  desc.add<int>("qualityForPrimary", 3);
  desc.add<bool>("usePVMET", true);
}

template <>
inline void PrimaryVertexSorter<std::vector<reco::RecoChargedRefCandidate>>::doConsumesForTiming(
    const edm::ParameterSet& iConfig) {
  tokenTrackTimeTag_ = consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("trackTimeTag"));
  tokenTrackTimeResoTag_ = consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("trackTimeResoTag"));
}

template <>
inline void PrimaryVertexSorter<std::vector<reco::PFCandidate>>::doConsumesForTiming(const edm::ParameterSet& iConfig) {
}

template <>
inline void PrimaryVertexSorter<std::vector<pat::PackedCandidate>>::doConsumesForTiming(
    const edm::ParameterSet& iConfig) {}

template <>
inline bool PrimaryVertexSorter<std::vector<reco::RecoChargedRefCandidate>>::needsProductsForTiming() {
  return true;
}

template <>
inline bool PrimaryVertexSorter<std::vector<reco::PFCandidate>>::needsProductsForTiming() {
  return false;
}

template <>
inline bool PrimaryVertexSorter<std::vector<pat::PackedCandidate>>::needsProductsForTiming() {
  return false;
}

template <>
inline std::pair<int, PrimaryVertexAssignment::Quality>
PrimaryVertexSorter<std::vector<reco::RecoChargedRefCandidate>>::runAlgo(const reco::VertexCollection& vertices,
                                                                         const reco::RecoChargedRefCandidate& pf,
                                                                         const edm::ValueMap<float>* trackTimeTag,
                                                                         const edm::ValueMap<float>* trackTimeResoTag,
                                                                         const edm::View<reco::Candidate>& jets,
                                                                         const TransientTrackBuilder& builder) {
  return assignmentAlgo_.chargedHadronVertex(vertices, pf, trackTimeTag, trackTimeResoTag, jets, builder);
}

template <>
inline std::pair<int, PrimaryVertexAssignment::Quality> PrimaryVertexSorter<std::vector<reco::PFCandidate>>::runAlgo(
    const reco::VertexCollection& vertices,
    const reco::PFCandidate& pf,
    const edm::ValueMap<float>* trackTimeTag,
    const edm::ValueMap<float>* trackTimeResoTag,
    const edm::View<reco::Candidate>& jets,
    const TransientTrackBuilder& builder) {
  return assignmentAlgo_.chargedHadronVertex(vertices, pf, jets, builder);
}

template <>
inline std::pair<int, PrimaryVertexAssignment::Quality> PrimaryVertexSorter<std::vector<pat::PackedCandidate>>::runAlgo(
    const reco::VertexCollection& vertices,
    const pat::PackedCandidate& pf,
    const edm::ValueMap<float>* trackTimeTag,
    const edm::ValueMap<float>* trackTimeResoTag,
    const edm::View<reco::Candidate>& jets,
    const TransientTrackBuilder& builder) {
  return assignmentAlgo_.chargedHadronVertex(vertices, pf, jets, builder);
}

#endif
