#include "FWCore/Framework/interface/one/EDAnalyzer.h"
#include "FWCore/Framework/interface/Run.h"
#include "FWCore/Framework/interface/Event.h"
#include "FWCore/Framework/interface/EventSetup.h"
#include "FWCore/Framework/interface/MakerMacros.h"
#include "FWCore/ParameterSet/interface/ParameterSet.h"
#include "FWCore/ServiceRegistry/interface/Service.h"
#include "FWCore/MessageLogger/interface/MessageLogger.h"
#include "FWCore/Utilities/interface/EDGetToken.h"
#include "FWCore/Utilities/interface/InputTag.h"
#include "FWCore/Utilities/interface/Exception.h"
#include "CommonTools/UtilAlgos/interface/TFileService.h"
#include "DataFormats/Common/interface/Handle.h"

#include "SimTracker/TrackTriggerAssociation/interface/StubAssociation.h"
#include "L1Trigger/TrackTrigger/interface/Setup.h"
#include "L1Trigger/TrackerTFP/interface/DataFormats.h"
#include "L1Trigger/TrackFindingTracklet/interface/ChannelAssignment.h"
#include "L1Trigger/TrackFindingTracklet/interface/Settings.h"

#include <TProfile.h>
#include <TProfile2D.h>
#include <TH1F.h>
#include <TH2F.h>

#include <vector>
#include <deque>
#include <set>
#include <cmath>
#include <numeric>
#include <sstream>

using namespace std;
using namespace edm;
using namespace trackerTFP;
using namespace tt;

namespace trklet {

  // stub resolution plots helper
  enum Resolution { Phi, Z, NumResolution };
  constexpr initializer_list<Resolution> AllResolution = {Phi, Z};
  constexpr auto NameResolution = {"Phi", "Z"};
  inline string name(Resolution r) { return string(*(NameResolution.begin() + r)); }

  /*! \class  trklet::AnalyzerTBout
   *  \brief  Class to analyze hardware like structured TTStub Collection generated by Tracklet
   *  \author Thomas Schuh
   *  \date   2021, Nov
   */
  class AnalyzerTBout : public one::EDAnalyzer<one::WatchRuns, one::SharedResources> {
  public:
    AnalyzerTBout(const ParameterSet& iConfig);
    void beginJob() override {}
    void beginRun(const Run& iEvent, const EventSetup& iSetup) override;
    void analyze(const Event& iEvent, const EventSetup& iSetup) override;
    void endRun(const Run& iEvent, const EventSetup& iSetup) override {}
    void endJob() override;

  private:
    //
    void formTracks(const StreamsTrack& streamsTrack,
                    const StreamsStub& streamsStubs,
                    vector<vector<TTStubRef>>& tracks,
                    int channel);
    //
    void associate(const vector<vector<TTStubRef>>& tracks, const StubAssociation* ass, set<TPPtr>& tps, int& sum) const;
    //
    void fill(const FrameTrack& frameTrack, const FrameStub& frameStub);

    // ED input token of dtc stubs
    EDGetTokenT<TTDTC> edGetTokenTTDTC_;
    // ED input token of stubs
    EDGetTokenT<StreamsStub> edGetTokenAcceptedStubs_;
    // ED input token of tracks
    EDGetTokenT<StreamsTrack> edGetTokenAcceptedTracks_;
    // ED input token of lost stubs
    EDGetTokenT<StreamsStub> edGetTokenLostStubs_;
    // ED input token of lost tracks
    EDGetTokenT<StreamsTrack> edGetTokenLostTracks_;
    // ED input token of TTStubRef to TPPtr association for tracking efficiency
    EDGetTokenT<StubAssociation> edGetTokenSelection_;
    // ED input token of TTStubRef to recontructable TPPtr association
    EDGetTokenT<StubAssociation> edGetTokenReconstructable_;
    // Setup token
    ESGetToken<Setup, SetupRcd> esGetTokenSetup_;
    // DataFormats token
    ESGetToken<DataFormats, DataFormatsRcd> esGetTokenDataFormats_;
    // ChannelAssignment token
    ESGetToken<ChannelAssignment, ChannelAssignmentRcd> esGetTokenChannelAssignment_;
    // stores, calculates and provides run-time constants
    const Setup* setup_ = nullptr;
    //
    const Settings settings_;
    // helper class to extract structured data from tt::Frames
    const DataFormats* dataFormats_ = nullptr;
    // helper class to assign tracklet track to channel
    const ChannelAssignment* channelAssignment_ = nullptr;
    // enables analyze of TPs
    bool useMCTruth_;
    //
    int nEvents_ = 0;
    //
    vector<deque<FrameStub>> regionStubs_;
    //
    int region_;
    //
    vector<int> nOverflows_;

    // Histograms

    TProfile* prof_;
    TProfile* profChannel_;
    TH1F* hisChannel_;
    vector<TH1F*> hisResolution_;
    vector<TProfile2D*> profResolution_;
    vector<TH1F*> hisResolutionMe_;
    vector<TH1F*> hisResolutionThey_;
    vector<TH2F*> his2Resolution_;

    // printout
    stringstream log_;
  };

  AnalyzerTBout::AnalyzerTBout(const ParameterSet& iConfig) : useMCTruth_(iConfig.getParameter<bool>("UseMCTruth")) {
    usesResource("TFileService");
    // book in- and output ED products
    const InputTag& inputTag = iConfig.getParameter<InputTag>("InputTagDTC");
    const string& label = iConfig.getParameter<string>("LabelTBout");
    const string& branchAcceptedStubs = iConfig.getParameter<string>("BranchAcceptedStubs");
    const string& branchAcceptedTracks = iConfig.getParameter<string>("BranchAcceptedTracks");
    const string& branchLostStubs = iConfig.getParameter<string>("BranchLostStubs");
    const string& branchLostTracks = iConfig.getParameter<string>("BranchLostTracks");
    edGetTokenTTDTC_ = consumes<TTDTC>(inputTag);
    edGetTokenAcceptedStubs_ = consumes<StreamsStub>(InputTag(label, branchAcceptedStubs));
    edGetTokenAcceptedTracks_ = consumes<StreamsTrack>(InputTag(label, branchAcceptedTracks));
    edGetTokenLostStubs_ = consumes<StreamsStub>(InputTag(label, branchLostStubs));
    edGetTokenLostTracks_ = consumes<StreamsTrack>(InputTag(label, branchLostTracks));
    if (useMCTruth_) {
      const auto& inputTagSelecttion = iConfig.getParameter<InputTag>("InputTagSelection");
      const auto& inputTagReconstructable = iConfig.getParameter<InputTag>("InputTagReconstructable");
      edGetTokenSelection_ = consumes<StubAssociation>(inputTagSelecttion);
      edGetTokenReconstructable_ = consumes<StubAssociation>(inputTagReconstructable);
    }
    // book ES products
    esGetTokenSetup_ = esConsumes<Setup, SetupRcd, Transition::BeginRun>();
    esGetTokenDataFormats_ = esConsumes<DataFormats, DataFormatsRcd, Transition::BeginRun>();
    esGetTokenChannelAssignment_ = esConsumes<ChannelAssignment, ChannelAssignmentRcd, Transition::BeginRun>();
    //
    nOverflows_ = vector<int>(2, 0);
    // log config
    log_.setf(ios::fixed, ios::floatfield);
    log_.precision(4);
  }

  void AnalyzerTBout::beginRun(const Run& iEvent, const EventSetup& iSetup) {
    // helper class to store configurations
    setup_ = &iSetup.getData(esGetTokenSetup_);
    // helper class to extract structured data from tt::Frames
    dataFormats_ = &iSetup.getData(esGetTokenDataFormats_);
    // helper class to assign tracklet track to channel
    channelAssignment_ = &iSetup.getData(esGetTokenChannelAssignment_);
    // book histograms
    Service<TFileService> fs;
    TFileDirectory dir;
    dir = fs->mkdir("TBout");
    prof_ = dir.make<TProfile>("Counts", ";", 9, 0.5, 9.5);
    prof_->GetXaxis()->SetBinLabel(1, "Stubs");
    prof_->GetXaxis()->SetBinLabel(2, "Tracks");
    prof_->GetXaxis()->SetBinLabel(3, "Lost Tracks");
    prof_->GetXaxis()->SetBinLabel(4, "Matched Tracks");
    prof_->GetXaxis()->SetBinLabel(5, "All Tracks");
    prof_->GetXaxis()->SetBinLabel(6, "Found TPs");
    prof_->GetXaxis()->SetBinLabel(7, "Found selected TPs");
    prof_->GetXaxis()->SetBinLabel(8, "Lost TPs");
    prof_->GetXaxis()->SetBinLabel(9, "All TPs");
    // channel occupancy
    constexpr int maxOcc = 180;
    const int numChannels = channelAssignment_->numChannelsStub() * setup_->numRegions();
    hisChannel_ = dir.make<TH1F>("His Channel Occupancy", ";", maxOcc, -.5, maxOcc - .5);
    profChannel_ = dir.make<TProfile>("Prof Channel Occupancy", ";", numChannels, -.5, numChannels - .5);
    // stub parameter resolutions
    constexpr int bins = 400;
    constexpr int binsHis = 100;
    constexpr double maxZ = 300.;
    constexpr double maxR = 120.;
    constexpr array<double, NumResolution> ranges{{.01, 5.}};
    hisResolution_.reserve(NumResolution);
    profResolution_.reserve(NumResolution);
    for (Resolution r : AllResolution) {
      hisResolution_.emplace_back(dir.make<TH1F>(("HisRes" + name(r)).c_str(), ";", binsHis, -ranges[r], ranges[r]));
      profResolution_.emplace_back(
          dir.make<TProfile2D>(("ProfRes" + name(r)).c_str(), ";;", bins, -maxZ, maxZ, bins, 0., maxR));
      hisResolutionMe_.emplace_back(
          dir.make<TH1F>(("HisResMe" + name(r)).c_str(), ";", binsHis, -ranges[r], ranges[r]));
      hisResolutionThey_.emplace_back(
          dir.make<TH1F>(("HisResThey" + name(r)).c_str(), ";", binsHis, -ranges[r], ranges[r]));
      his2Resolution_.emplace_back(
          dir.make<TH2F>(("His2" + name(r)).c_str(), ";;", bins, -ranges[r], ranges[r], bins, -ranges[r], ranges[r]));
    }
    regionStubs_ = vector<deque<FrameStub>>(setup_->numRegions());
  }

  void AnalyzerTBout::analyze(const Event& iEvent, const EventSetup& iSetup) {
    // read in TTDTC
    Handle<TTDTC> handleTTDTC;
    iEvent.getByToken<TTDTC>(edGetTokenTTDTC_, handleTTDTC);
    const TTDTC& ttDTC = *handleTTDTC;
    for (deque<FrameStub>& region : regionStubs_)
      region.clear();
    for (int r : ttDTC.tfpRegions()) {
      for (int c : ttDTC.tfpChannels()) {
        const StreamStub& s = ttDTC.stream(r, c);
        copy_if(
            s.begin(), s.end(), back_inserter(regionStubs_[r]), [](const FrameStub& f) { return f.first.isNonnull(); });
      }
    }
    // read in TBout products
    Handle<StreamsStub> handleAcceptedStubs;
    iEvent.getByToken<StreamsStub>(edGetTokenAcceptedStubs_, handleAcceptedStubs);
    const StreamsStub& acceptedStubs = *handleAcceptedStubs;
    Handle<StreamsTrack> handleAcceptedTracks;
    iEvent.getByToken<StreamsTrack>(edGetTokenAcceptedTracks_, handleAcceptedTracks);
    const StreamsTrack& acceptedTracks = *handleAcceptedTracks;
    Handle<StreamsStub> handleLostStubs;
    iEvent.getByToken<StreamsStub>(edGetTokenLostStubs_, handleLostStubs);
    const StreamsStub& lostStubs = *handleLostStubs;
    Handle<StreamsTrack> handleLostTracks;
    iEvent.getByToken<StreamsTrack>(edGetTokenLostTracks_, handleLostTracks);
    const StreamsTrack& lostTracks = *handleLostTracks;
    // read in MCTruth
    const StubAssociation* selection = nullptr;
    const StubAssociation* reconstructable = nullptr;
    if (useMCTruth_) {
      Handle<StubAssociation> handleSelection;
      iEvent.getByToken<StubAssociation>(edGetTokenSelection_, handleSelection);
      selection = handleSelection.product();
      prof_->Fill(9, selection->numTPs());
      Handle<StubAssociation> handleReconstructable;
      iEvent.getByToken<StubAssociation>(edGetTokenReconstructable_, handleReconstructable);
      reconstructable = handleReconstructable.product();
    }
    // analyze ht products and associate found tracks with reconstrucable TrackingParticles
    set<TPPtr> tpPtrs;
    set<TPPtr> tpPtrsSelection;
    set<TPPtr> tpPtrsLost;
    int allMatched(0);
    int allTracks(0);
    for (region_ = 0; region_ < setup_->numRegions(); region_++) {
      const int offset = region_ * channelAssignment_->numChannelsTrack();
      int nStubs(0);
      int nTracks(0);
      int nLost(0);
      for (int channel = 0; channel < channelAssignment_->numChannelsTrack(); channel++) {
        vector<vector<TTStubRef>> tracks;
        formTracks(acceptedTracks, acceptedStubs, tracks, offset + channel);
        vector<vector<TTStubRef>> lost;
        formTracks(lostTracks, lostStubs, lost, offset + channel);
        nTracks += tracks.size();
        nStubs +=
            accumulate(tracks.begin(), tracks.end(), 0, [](int sum, const auto& v) { return sum + (int)v.size(); });
        nLost += lost.size();
        allTracks += tracks.size();
        if (!useMCTruth_)
          continue;
        int tmp(0);
        associate(tracks, selection, tpPtrsSelection, tmp);
        associate(lost, selection, tpPtrsLost, tmp);
        associate(tracks, reconstructable, tpPtrs, allMatched);
      }
      prof_->Fill(1, nStubs);
      prof_->Fill(2, nTracks);
      prof_->Fill(3, nLost);
    }
    vector<TPPtr> recovered;
    recovered.reserve(tpPtrsLost.size());
    set_intersection(tpPtrsLost.begin(), tpPtrsLost.end(), tpPtrs.begin(), tpPtrs.end(), back_inserter(recovered));
    for (const TPPtr& tpPtr : recovered)
      tpPtrsLost.erase(tpPtr);
    prof_->Fill(4, allMatched);
    prof_->Fill(5, allTracks);
    prof_->Fill(6, tpPtrs.size());
    prof_->Fill(7, tpPtrsSelection.size());
    prof_->Fill(8, tpPtrsLost.size());
    nEvents_++;
  }

  void AnalyzerTBout::endJob() {
    if (nEvents_ == 0)
      return;
    // printout TBout summary
    const double totalTPs = prof_->GetBinContent(9);
    const double numStubs = prof_->GetBinContent(1);
    const double numTracks = prof_->GetBinContent(2);
    const double numTracksLost = prof_->GetBinContent(3);
    const double totalTracks = prof_->GetBinContent(5);
    const double numTracksMatched = prof_->GetBinContent(4);
    const double numTPsAll = prof_->GetBinContent(6);
    const double numTPsEff = prof_->GetBinContent(7);
    const double numTPsLost = prof_->GetBinContent(8);
    const double errStubs = prof_->GetBinError(1);
    const double errTracks = prof_->GetBinError(2);
    const double errTracksLost = prof_->GetBinError(3);
    const double fracFake = (totalTracks - numTracksMatched) / totalTracks;
    const double fracDup = (numTracksMatched - numTPsAll) / totalTracks;
    const double eff = numTPsEff / totalTPs;
    const double errEff = sqrt(eff * (1. - eff) / totalTPs / nEvents_);
    const double effLoss = numTPsLost / totalTPs;
    const double errEffLoss = sqrt(effLoss * (1. - effLoss) / totalTPs / nEvents_);
    const vector<double> nums = {numStubs, numTracks, numTracksLost};
    const vector<double> errs = {errStubs, errTracks, errTracksLost};
    const int wNums = ceil(log10(*max_element(nums.begin(), nums.end()))) + 5;
    const int wErrs = ceil(log10(*max_element(errs.begin(), errs.end()))) + 5;
    log_ << "                        TBout SUMMARY                        " << endl;
    log_ << "number of stubs       per TFP = " << setw(wNums) << numStubs << " +- " << setw(wErrs) << errStubs << endl;
    log_ << "number of tracks      per TFP = " << setw(wNums) << numTracks << " +- " << setw(wErrs) << errTracks
         << endl;
    log_ << "number of lost tracks per TFP = " << setw(wNums) << numTracksLost << " +- " << setw(wErrs) << errTracksLost
         << endl;
    log_ << "     max  tracking efficiency = " << setw(wNums) << eff << " +- " << setw(wErrs) << errEff << endl;
    log_ << "     lost tracking efficiency = " << setw(wNums) << effLoss << " +- " << setw(wErrs) << errEffLoss << endl;
    log_ << "                    fake rate = " << setw(wNums) << fracFake << endl;
    log_ << "               duplicate rate = " << setw(wNums) << fracDup << endl;
    log_ << "number of overflowed phi residuals: " << nOverflows_[0] << " and z: " << nOverflows_[1] << endl;
    log_ << "=============================================================";
    LogPrint("L1Trigger/TrackerTFP") << log_.str();
  }

  //
  void AnalyzerTBout::formTracks(const StreamsTrack& streamsTrack,
                                 const StreamsStub& streamsStubs,
                                 vector<vector<TTStubRef>>& tracks,
                                 int channel) {
    const int seedType = channel % channelAssignment_->numChannelsTrack();
    const int offset = channelAssignment_->offsetStub(channel);
    const StreamTrack& streamTrack = streamsTrack[channel];
    const int numTracks = accumulate(streamTrack.begin(), streamTrack.end(), 0, [](int sum, const FrameTrack& frame) {
      return sum + (frame.first.isNonnull() ? 1 : 0);
    });
    tracks.reserve(numTracks);
    for (int frame = 0; frame < (int)streamTrack.size(); frame++) {
      const FrameTrack& frameTrack = streamTrack[frame];
      if (frameTrack.first.isNull())
        continue;
      vector<TTStubRef> ttStubRefs;
      ttStubRefs.reserve(channelAssignment_->numProjectionLayers(seedType) +
                         channelAssignment_->seedingLayers(seedType).size());
      for (int layer = 0; layer < channelAssignment_->numProjectionLayers(seedType); layer++) {
        const FrameStub& stub = streamsStubs[offset + layer][frame];
        if (stub.first.isNonnull()) {
          this->fill(frameTrack, stub);
          ttStubRefs.push_back(stub.first);
        }
      }
      for (const TTStubRef& ttStubRef : frameTrack.first->getStubRefs()) {
        int layer(0);
        if (!channelAssignment_->layerId(seedType, ttStubRef, layer))
          ttStubRefs.push_back(ttStubRef);
      }
      tracks.push_back(ttStubRefs);
    }
  }

  //
  void AnalyzerTBout::associate(const vector<vector<TTStubRef>>& tracks,
                                const StubAssociation* ass,
                                set<TPPtr>& tps,
                                int& sum) const {
    for (const vector<TTStubRef>& ttStubRefs : tracks) {
      const vector<TPPtr>& tpPtrs = ass->associate(ttStubRefs);
      if (tpPtrs.empty())
        continue;
      sum++;
      copy(tpPtrs.begin(), tpPtrs.end(), inserter(tps, tps.begin()));
    }
  }

  //
  void AnalyzerTBout::fill(const FrameTrack& frameTrack, const FrameStub& frameStub) {
    // get dtc stub frame
    const deque<FrameStub>& region = regionStubs_[region_];
    const auto it =
        find_if(region.begin(), region.end(), [&frameStub](const FrameStub& f) { return f.first == frameStub.first; });
    if (it == region.end())
      throw cms::Exception("LgociError.") << "Stub on track was not in DTC collection.";
    const GlobalPoint ttPos = setup_->stubPos(frameStub.first);
    const GlobalPoint pos = setup_->stubPos(true, *it, region_);
    static constexpr int widthPhi = 12;
    static constexpr int widthZ = 9;
    static constexpr int widthR = 7;
    const bool barrel = setup_->barrel(frameStub.first);
    const int layerIdTracklet = setup_->trackletLayerId(frameStub.first);
    static const double baseR = settings_.kz();
    const double basePhi = barrel ? settings_.kphi1() : settings_.kphi(layerIdTracklet);
    const double baseZ = settings_.kz(layerIdTracklet);
    static const double baseInvR = settings_.kphi1() / settings_.kr() * pow(2, settings_.rinv_shift());
    static const double basePhi0 = settings_.kphi1() * pow(2, settings_.phi0_shift());
    static const double baseZ0 = settings_.kz() * pow(2, settings_.z0_shift());
    static const double baseTanL = settings_.kz() / settings_.kr() * pow(2, settings_.t_shift());
    const int widthRZ = barrel ? widthZ : widthR;
    const double baseRZ = barrel ? baseZ : baseR;
    // calc residuals
    const double rInv = (frameTrack.first->rInv() / baseInvR + .5) * baseInvR;
    const double phi0 = (frameTrack.first->phi() / basePhi0 + .5) * basePhi0;
    const double z0 = (frameTrack.first->z0() / baseZ0 + .5) * baseZ0;
    const double tanL = (frameTrack.first->tanL() / baseTanL + .5) * baseTanL;
    const double phi = deltaPhi(pos.phi() - (phi0 - rInv * pos.perp() / 2.));
    const double r = pos.perp() - (pos.z() - z0) / tanL;
    const double z = pos.z() - (z0 + tanL * pos.perp());
    const double rz = barrel ? z : r;
    const int phii = floor(phi / basePhi);
    const int rzi = floor(rz / baseRZ);
    const double phid = (phii + .5) * basePhi;
    const double rzd = (rzi + .5) * baseRZ;
    // parse residuals
    TTBV hw(frameStub.second);
    const TTBV hwRZ(hw, widthRZ, 0, true);
    hw >>= widthRZ;
    const TTBV hwPhi(hw, widthPhi, 0, true);
    bool overflowPhi = abs(phii - hwPhi.val()) > pow(2, widthPhi) * 7. / 8.;
    bool overflowZ = abs(rzi - hwRZ.val()) > pow(2, widthRZ) * 7. / 8.;
    const double hwPhid = hwPhi.val(basePhi);
    const double hwRZd = hwRZ.val(baseRZ);
    const vector<double> resolutions = {phid - hwPhid, rzd - hwRZd};
    const vector<bool> overflows = {overflowPhi, overflowZ};
    for (Resolution r : AllResolution) {
      if (overflows[r]) {
        cout << rzi << " " << hwRZ.val() << " " << barrel << " " << setup_->psModule(frameStub.first) << " "
             << frameStub.second << endl;
        nOverflows_[r]++;
        continue;
      }
      hisResolution_[r]->Fill(resolutions[r]);
      profResolution_[r]->Fill(ttPos.z(), ttPos.perp(), abs(resolutions[r]));
    }
    hisResolutionMe_[0]->Fill(phid);
    hisResolutionMe_[1]->Fill(rzd);
    hisResolutionThey_[0]->Fill(hwPhid);
    hisResolutionThey_[1]->Fill(hwRZd);
    his2Resolution_[0]->Fill(phid, hwPhid);
    his2Resolution_[1]->Fill(rzd, hwRZd);
  }

}  // namespace trklet

DEFINE_FWK_MODULE(trklet::AnalyzerTBout);
