#include <memory>

// user include files
#include "FWCore/Framework/interface/Frameworkfwd.h"
#include "FWCore/Framework/interface/stream/EDProducer.h"
#include "DataFormats/TrackReco/interface/Track.h"
#include "FWCore/Framework/interface/Event.h"
#include "FWCore/Framework/interface/MakerMacros.h"
#include "TrackingTools/PatternTools/interface/TrajTrackAssociation.h"

class dso_hidden QualityFilter final : public edm::stream::EDProducer<> {
public:
  explicit QualityFilter(const edm::ParameterSet&);
  ~QualityFilter() override;

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

  // ----------member data ---------------------------
private:
  edm::EDGetTokenT<std::vector<Trajectory> > trajTag;
  edm::EDGetTokenT<TrajTrackAssociationCollection> tassTag;
  reco::TrackBase::TrackQuality trackQuality_;
  bool copyExtras_;
};

#include "DataFormats/TrackReco/interface/TrackExtra.h"
#include "DataFormats/TrackReco/interface/TrackExtraFwd.h"
#include "FWCore/Utilities/interface/InputTag.h"
#include "TrackingTools/TransientTrackingRecHit/interface/TransientTrackingRecHit.h"
#include "TrackingTools/PatternTools/interface/Trajectory.h"
#include "TrackingTools/PatternTools/interface/TrajTrackAssociation.h"

//
// class decleration
//

using namespace edm;
using namespace reco;
using namespace std;

QualityFilter::QualityFilter(const edm::ParameterSet& iConfig) {
  copyExtras_ = iConfig.getUntrackedParameter<bool>("copyExtras", false);

  produces<reco::TrackCollection>();
  if (copyExtras_) {
    produces<TrackingRecHitCollection>();
    produces<reco::TrackExtraCollection>();
  }
  produces<std::vector<Trajectory> >();
  produces<TrajTrackAssociationCollection>();

  trajTag = consumes<std::vector<Trajectory> >(iConfig.getParameter<edm::InputTag>("recTracks"));
  tassTag = consumes<TrajTrackAssociationCollection>(iConfig.getParameter<edm::InputTag>("recTracks"));
  trackQuality_ = TrackBase::qualityByName(iConfig.getParameter<std::string>("TrackQuality"));
}

QualityFilter::~QualityFilter() {
  // do anything here that needs to be done at desctruction time
  // (e.g. close files, deallocate resources etc.)
}

//
// member functions
//

// ------------ method called to produce the data  ------------
void QualityFilter::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
  unique_ptr<TrackCollection> selTracks(new TrackCollection);
  unique_ptr<TrackingRecHitCollection> selHits(copyExtras_ ? new TrackingRecHitCollection() : nullptr);
  unique_ptr<TrackExtraCollection> selTrackExtras(copyExtras_ ? new TrackExtraCollection() : nullptr);
  unique_ptr<vector<Trajectory> > outputTJ(new vector<Trajectory>);
  unique_ptr<TrajTrackAssociationCollection> trajTrackMap(new TrajTrackAssociationCollection());

  TrackExtraRefProd rTrackExtras;
  TrackingRecHitRefProd rHits;
  if (copyExtras_) {
    rTrackExtras = iEvent.getRefBeforePut<TrackExtraCollection>();
    rHits = iEvent.getRefBeforePut<TrackingRecHitCollection>();
  }

  Handle<std::vector<Trajectory> > TrajectoryCollection;
  Handle<TrajTrackAssociationCollection> assoMap;

  iEvent.getByToken(trajTag, TrajectoryCollection);
  iEvent.getByToken(tassTag, assoMap);

  TrajTrackAssociationCollection::const_iterator it = assoMap->begin();
  TrajTrackAssociationCollection::const_iterator lastAssoc = assoMap->end();
  for (; it != lastAssoc; ++it) {
    const Ref<vector<Trajectory> > traj = it->key;
    const reco::TrackRef itc = it->val;
    bool goodTk = (itc->quality(trackQuality_));

    if (goodTk) {
      auto const& track = (*itc);
      //tracks and trajectories
      selTracks->push_back(track);
      outputTJ->push_back(*traj);
      if (copyExtras_) {
        //TRACKING HITS
        trackingRecHit_iterator irhit = (*itc).recHitsBegin();
        trackingRecHit_iterator lasthit = (*itc).recHitsEnd();
        for (; irhit != lasthit; ++irhit) {
          selHits->push_back((*irhit)->clone());
        }
      }
    }
  }

  unsigned nTracks = selTracks->size();
  if (copyExtras_) {
    //PUT TRACKING HITS IN THE EVENT
    OrphanHandle<TrackingRecHitCollection> theRecoHits = iEvent.put(std::move(selHits));
    edm::RefProd<TrackingRecHitCollection> theRecoHitsProd(theRecoHits);

    //PUT TRACK EXTRA IN THE EVENT
    selTrackExtras->reserve(nTracks);
    unsigned hits = 0;

    for (unsigned index = 0; index < nTracks; ++index) {
      auto const& aTrack = (*selTracks)[index];
      selTrackExtras->emplace_back(aTrack.outerPosition(),
                                   aTrack.outerMomentum(),
                                   aTrack.outerOk(),
                                   aTrack.innerPosition(),
                                   aTrack.innerMomentum(),
                                   aTrack.innerOk(),
                                   aTrack.outerStateCovariance(),
                                   aTrack.outerDetId(),
                                   aTrack.innerStateCovariance(),
                                   aTrack.innerDetId(),
                                   aTrack.seedDirection(),
                                   aTrack.seedRef());

      auto& aTrackExtra = selTrackExtras->back();
      //unsigned nHits = aTrack.numberOfValidHits();
      auto nHits = aTrack.recHitsSize();
      aTrackExtra.setHits(theRecoHitsProd, hits, nHits);
      hits += nHits;
      selTrackExtras->push_back(aTrackExtra);
    }

    //CORRECT REF TO TRACK
    OrphanHandle<TrackExtraCollection> theRecoTrackExtras = iEvent.put(std::move(selTrackExtras));
    for (unsigned index = 0; index < nTracks; ++index) {
      const reco::TrackExtraRef theTrackExtraRef(theRecoTrackExtras, index);
      (*selTracks)[index].setExtra(theTrackExtraRef);
    }
  }  // END IF COPY EXTRAS

  //TRACKS AND TRAJECTORIES
  OrphanHandle<TrackCollection> theRecoTracks = iEvent.put(std::move(selTracks));
  OrphanHandle<vector<Trajectory> > theRecoTrajectories = iEvent.put(std::move(outputTJ));

  //TRACKS<->TRAJECTORIES MAP
  nTracks = theRecoTracks->size();
  for (unsigned index = 0; index < nTracks; ++index) {
    Ref<vector<Trajectory> > trajRef(theRecoTrajectories, index);
    Ref<TrackCollection> tkRef(theRecoTracks, index);
    trajTrackMap->insert(trajRef, tkRef);
  }
  //MAP IN THE EVENT
  iEvent.put(std::move(trajTrackMap));
}

// ------------ method called once each job just after ending the event loop  ------------
void QualityFilter::endJob() {}

DEFINE_FWK_MODULE(QualityFilter);
