#include <alpaka/alpaka.hpp>

#include "DataFormats/Portable/interface/Product.h"
#include "FWCore/Framework/interface/Event.h"
#include "FWCore/Framework/interface/EventSetup.h"
#include "FWCore/Framework/interface/Frameworkfwd.h"
#include "FWCore/Framework/interface/stream/EDProducer.h"
#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
#include "FWCore/ParameterSet/interface/ParameterSet.h"
#include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
#include "FWCore/Utilities/interface/EDGetToken.h"
#include "FWCore/Utilities/interface/InputTag.h"
#include "FWCore/Utilities/interface/StreamID.h"
#include "HeterogeneousCore/AlpakaCore/interface/ScopedContext.h"
#include "HeterogeneousCore/AlpakaInterface/interface/config.h"

#include "RecoTracker/LST/interface/LSTPixelSeedInput.h"
#include "RecoTracker/LST/interface/LSTPhase2OTHitsInput.h"
#include "RecoTracker/LST/interface/LSTOutput.h"

#include "HeterogeneousCore/SonicTriton/interface/TritonEDProducer.h"

#ifdef ALPAKA_ACC_GPU_CUDA_ENABLED
#include "SDL/LST.h"
#endif  // ALPAKA_ACC_GPU_CUDA_ENABLED

namespace ALPAKA_ACCELERATOR_NAMESPACE {

  class LSTProducer : public TritonEDProducer<> {
  public:
    LSTProducer(edm::ParameterSet& config)
        : lstPixelSeedInputToken_{consumes<LSTPixelSeedInput>(config.getParameter<edm::InputTag>("pixelSeedInput"))},
          lstPhase2OTHitsInputToken_{
              consumes<LSTPhase2OTHitsInput>(config.getParameter<edm::InputTag>("phase2OTHitsInput"))},
          verbose_(config.getParameter<int>("verbose")),
          lstOutputToken_{produces<LSTOutput>()} {}

    //void acquire(edm::Event const& event, edm::EventSetup const& setup, edm::WaitingTaskWithArenaHolder task) override {
    void acquire(edm::Event const& iEvent, edm::EventSetup const& iSetup, edm::WaitingTaskWithArenaHolder task, Input& iInput) {
      // create a context based on the EDM stream number
      cms::alpakatools::ScopedContextAcquire<Queue> ctx{iEvent.streamID(), std::move(task)};

      // Inputs
      auto const& pixelSeeds = iEvent.get(lstPixelSeedInputToken_);
      auto const& phase2OTHits = iEvent.get(lstPhase2OTHitsInputToken_);


      ////auto& input = iInput.at("input");
      ////auto pfdata = input.allocate<float>();
      ////auto& vpfdata = (*pfdata)[0];
      ////for (const auto& pixelSeed : pixelSeeds) {
        ////vpfdata.push_back(pixelSeed.dx());
      ////}

#ifdef ALPAKA_ACC_GPU_CUDA_ENABLED
      lst_.eventSetup();
      lst_.run(ctx.queue().getNativeHandle(),
               verbose_,
               pixelSeeds.px(),
               pixelSeeds.py(),
               pixelSeeds.pz(),
               pixelSeeds.dxy(),
               pixelSeeds.dz(),
               pixelSeeds.ptErr(),
               pixelSeeds.etaErr(),
               pixelSeeds.stateTrajGlbX(),
               pixelSeeds.stateTrajGlbY(),
               pixelSeeds.stateTrajGlbZ(),
               pixelSeeds.stateTrajGlbPx(),
               pixelSeeds.stateTrajGlbPy(),
               pixelSeeds.stateTrajGlbPz(),
               pixelSeeds.q(),
               pixelSeeds.hitIdx(),
               phase2OTHits.detId(),
               phase2OTHits.x(),
               phase2OTHits.y(),
               phase2OTHits.z());
#endif  // ALPAKA_ACC_GPU_CUDA_ENABLED
    }

    void produce(edm::Event& iEvent, edm::EventSetup const& iSetup, Output const& iOutput) {
      // Output
      LSTOutput lstOutput;
      lstOutput.setLSTOutputTraits(lst_.hits(), lst_.len(), lst_.seedIdx(), lst_.trackCandidateType());

      iEvent.emplace(lstOutputToken_, std::move(lstOutput));
    }

    static void fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
      edm::ParameterSetDescription desc;
      desc.add<edm::InputTag>("pixelSeedInput", edm::InputTag{"lstPixelSeedInputProducer"});
      desc.add<edm::InputTag>("phase2OTHitsInput", edm::InputTag{"lstPhase2OTHitsInputProducer"});
      desc.add<int>("verbose", 0);
      descriptions.addWithDefaultLabel(desc);
    }

  private:
    edm::EDGetTokenT<LSTPixelSeedInput> lstPixelSeedInputToken_;
    edm::EDGetTokenT<LSTPhase2OTHitsInput> lstPhase2OTHitsInputToken_;
    const int verbose_;
    edm::EDPutTokenT<LSTOutput> lstOutputToken_;

#ifdef ALPAKA_ACC_GPU_CUDA_ENABLED
    SDL::LST lst_;
#endif  // ALPAKA_ACC_GPU_CUDA_ENABLED
  };

}  // namespace ALPAKA_ACCELERATOR_NAMESPACE

#include "HeterogeneousCore/AlpakaCore/interface/MakerMacros.h"
DEFINE_FWK_ALPAKA_MODULE(LSTProducer);