text: /****************************************************************************
   * Authors:
   *   Jan Kašpar
   ****************************************************************************/
  
  #include "FWCore/Framework/interface/stream/EDProducer.h"
  #include "FWCore/Framework/interface/Event.h"
  #include "FWCore/ParameterSet/interface/ParameterSet.h"
  #include "FWCore/Framework/interface/EventSetup.h"
  #include "FWCore/Framework/interface/ESHandle.h"
  #include "FWCore/Framework/interface/MakerMacros.h"
  #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
  #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
  #include "FWCore/ServiceRegistry/interface/Service.h"
  #include "FWCore/Utilities/interface/RandomNumberGenerator.h"
  
  #include "SimDataFormats/GeneratorProducts/interface/HepMCProduct.h"
  
  #include "DataFormats/HepMCCandidate/interface/GenParticle.h"
  
  #include "CondFormats/CTPPSReadoutObjects/interface/CTPPSBeamParameters.h"
  #include "CondFormats/DataRecord/interface/CTPPSBeamParametersRcd.h"
  
  #include <CLHEP/Random/RandGauss.h>
  
  //----------------------------------------------------------------------------------------------------
  
  class BeamDivergenceVtxGenerator : public edm::stream::EDProducer<>
  {
    public:
      explicit BeamDivergenceVtxGenerator(const edm::ParameterSet&);
      ~BeamDivergenceVtxGenerator() override = default;
  
      static void fillDescriptions(edm::ConfigurationDescriptions&);
  
      void produce(edm::Event&, const edm::EventSetup&) override;
  
    private:
      edm::EDGetTokenT<edm::HepMCProduct> sourceToken_;
      std::vector<edm::EDGetTokenT<reco::GenParticleCollection>> genParticleTokens_;
  
      bool simulateVertex_;
      bool simulateBeamDivergence_;
  
      struct SmearingParameters
      {
        double vtx_x, vtx_y, vtx_z; // cm
        double bd_x_45, bd_y_45, bd_x_56, bd_y_56; // rad
      };
  
      template <typename T>
      static HepMC::FourVector smearMomentum(const T &mom, const SmearingParameters &sp);
  
      void applySmearingHepMC(const SmearingParameters &sp, HepMC::GenEvent *genEvt);
  
      void addSmearedGenParticle(const reco::GenParticle &gp, const SmearingParameters &sp, HepMC::GenEvent *genEvt);
  };
  
  //----------------------------------------------------------------------------------------------------
  
  BeamDivergenceVtxGenerator::BeamDivergenceVtxGenerator(const edm::ParameterSet& iConfig) :
    simulateVertex_        (iConfig.getParameter<bool>("simulateVertex")),
    simulateBeamDivergence_(iConfig.getParameter<bool>("simulateBeamDivergence"))
  {
    const edm::InputTag tagSrcHepMC = iConfig.getParameter<edm::InputTag>("src");
    if (tagSrcHepMC.label() != "")
      sourceToken_ = consumes<edm::HepMCProduct>(tagSrcHepMC);
  
    for (const auto &it : iConfig.getParameter<std::vector<edm::InputTag>>("srcGenParticle"))
      genParticleTokens_.push_back(consumes<reco::GenParticleCollection>(it));
  
    edm::Service<edm::RandomNumberGenerator> rng;
    if (!rng.isAvailable())
      throw cms::Exception("Configuration")
        << "The BeamDivergenceVtxGenerator requires the RandomNumberGeneratorService\n"
           "which is not present in the configuration file. \n"
           "You must add the service\n"
           "in the configuration file or remove the modules that require it.";
  
    produces<edm::HepMCProduct>();
  }
  
  //----------------------------------------------------------------------------------------------------
  
  void BeamDivergenceVtxGenerator::fillDescriptions(edm::ConfigurationDescriptions& descriptions)
  {
    edm::ParameterSetDescription desc;
  
    desc.add<edm::InputTag>("src", edm::InputTag("generator", "unsmeared"))
      ->setComment("input collection in HepMC format");
  
    desc.add<std::vector<edm::InputTag>>("srcGenParticle", std::vector<edm::InputTag>())
      ->setComment("input collections in GenParticle format");
  
    desc.add<bool>("simulateBeamDivergence", true)->setComment("account for the beam angular divergence?");
    desc.add<bool>("simulateVertex", true)->setComment("account for the vertex transverse smearing?");
  
    descriptions.add("beamDivergenceVtxGenerator", desc);
  }
  
  //----------------------------------------------------------------------------------------------------
  
  void BeamDivergenceVtxGenerator::produce(edm::Event& iEvent, const edm::EventSetup &iSetup)
  {
    // get random engine
    edm::Service<edm::RandomNumberGenerator> rng;
    CLHEP::HepRandomEngine* rnd = &(rng->getEngine(iEvent.streamID()));
  
    // get conditions
    edm::ESHandle<CTPPSBeamParameters> hBeamParameters;
    iSetup.get<CTPPSBeamParametersRcd>().get(hBeamParameters);
  
    // get HepMC input (if given)
    HepMC::GenEvent* genEvt;
    if (sourceToken_.isUninitialized())
    {
      genEvt = new HepMC::GenEvent();
    } else {
      edm::Handle<edm::HepMCProduct> hepUnsmearedMCEvt;
      iEvent.getByToken(sourceToken_, hepUnsmearedMCEvt);
  
      genEvt = new HepMC::GenEvent(*hepUnsmearedMCEvt->GetEvent());
    }
  
    // prepare output
    std::unique_ptr<edm::HepMCProduct> output(new edm::HepMCProduct(genEvt));
  
    // generate smearing parameters
    SmearingParameters sp;
  
    if (simulateVertex_)
    {
      // NB: the separtion between effective offsets in LHC sectors 45 and 56 cannot be applied, thus the values for 45 are used
      sp.vtx_x = hBeamParameters->getVtxOffsetX45() + CLHEP::RandGauss::shoot(rnd) * hBeamParameters->getVtxStddevX();
      sp.vtx_y = hBeamParameters->getVtxOffsetY45() + CLHEP::RandGauss::shoot(rnd) * hBeamParameters->getVtxStddevY();
      sp.vtx_z = hBeamParameters->getVtxOffsetZ45() + CLHEP::RandGauss::shoot(rnd) * hBeamParameters->getVtxStddevZ();
    }
  
    if (simulateBeamDivergence_)
    {
      sp.bd_x_45 = CLHEP::RandGauss::shoot(rnd) * hBeamParameters->getBeamDivergenceX45();
      sp.bd_x_56 = CLHEP::RandGauss::shoot(rnd) * hBeamParameters->getBeamDivergenceX56();
      sp.bd_y_45 = CLHEP::RandGauss::shoot(rnd) * hBeamParameters->getBeamDivergenceY45();
      sp.bd_y_56 = CLHEP::RandGauss::shoot(rnd) * hBeamParameters->getBeamDivergenceY56();
    }
  
    // apply smearing
    applySmearingHepMC(sp, genEvt);
  
    for (const auto &token : genParticleTokens_)
    {
      edm::Handle<reco::GenParticleCollection> hGPCollection;
      iEvent.getByToken(token, hGPCollection);
  
      for (const auto &gp : *hGPCollection)
        addSmearedGenParticle(gp, sp, genEvt);
    }
  
    // save output
    iEvent.put(std::move(output));
  }
  
  //----------------------------------------------------------------------------------------------------
  
  template <typename T>
  HepMC::FourVector BeamDivergenceVtxGenerator::smearMomentum(const T &mom, const SmearingParameters &sp)
  {
    // TODO: this is an oversimplified implemetation
    // the TOTEM smearing module should be taken as reference
  
    double th_x = mom.x() / mom.z();
    double th_y = mom.y() / mom.z();
  
    if (mom.z() > 0.)
    {
      th_x += sp.bd_x_45;
      th_y += sp.bd_y_45;
    } else {
      th_x += sp.bd_x_56;
      th_y += sp.bd_y_56;
    }
  
    // calculate consistent p_z component
    const double sign = (mom.z() > 0.) ? 1. : -1.;
    const double p = sqrt(mom.x()*mom.x() + mom.y()*mom.y() + mom.z()*mom.z());
    const double p_z = sign * p / sqrt(1. + th_x*th_x + th_y*th_y);
  
    return HepMC::FourVector(p_z * th_x, p_z * th_y, p_z, mom.e());
  }
  
  //----------------------------------------------------------------------------------------------------
  
  void BeamDivergenceVtxGenerator::applySmearingHepMC(const SmearingParameters &sp, HepMC::GenEvent *genEvt)
  {
    if (simulateVertex_)
    {
      for (HepMC::GenEvent::vertex_iterator vit = genEvt->vertices_begin(); vit != genEvt->vertices_end(); ++vit)
      {
        const auto &pos = (*vit)->position();
        (*vit)->set_position(HepMC::FourVector(
          pos.x() + sp.vtx_x*1E1, // conversion: cm to mm
          pos.y() + sp.vtx_y*1E1,
          pos.z() + sp.vtx_z*1E1,
          pos.t()
        ));
      }
    }
  
    if (simulateBeamDivergence_)
    {
      for (HepMC::GenEvent::particle_iterator part = genEvt->particles_begin(); part != genEvt->particles_end(); ++part)
        (*part)->set_momentum(smearMomentum((*part)->momentum(), sp));
    }
  }
  
  //----------------------------------------------------------------------------------------------------
  
  void BeamDivergenceVtxGenerator::addSmearedGenParticle(const reco::GenParticle &gp, const SmearingParameters &sp, HepMC::GenEvent *genEvt)
  {
    // add vertex of the particle
    HepMC::GenVertex *vtx = new HepMC::GenVertex(HepMC::FourVector(
      (gp.vx() + sp.vtx_x)*1E1, // conversion: cm to mm
      (gp.vy() + sp.vtx_y)*1E1,
      (gp.vz() + sp.vtx_z)*1E1,
      0.
    ));
    genEvt->add_vertex(vtx);
  
    // add the particle itself
    HepMC::GenParticle* particle = new HepMC::GenParticle(smearMomentum(gp.p4(), sp), gp.pdgId(), gp.status());
    vtx->add_particle_out(particle);
  }
  
  //----------------------------------------------------------------------------------------------------
  
  DEFINE_FWK_MODULE(BeamDivergenceVtxGenerator);
  