// -*- C++ -*-
#include "Rivet/Analysis.hh"
#include "Rivet/Projections/FinalState.hh"
#include "Rivet/Projections/DileptonFinder.hh"

namespace Rivet {


  /// @brief Forward Z production cross-section in pp collisions at 13 TeV
  class LHCB_2021_I1990313 : public Analysis {
  public:

    /// Constructor
    RIVET_DEFAULT_ANALYSIS_CTOR(LHCB_2021_I1990313);


    /// @name Analysis methods
    ///@{

    /// Book histograms and initialise projections before the run
    void init() {

      // Initialise and register projections
      DileptonFinder zmumufinder(91.2*GeV, 0.1, Cuts::absetaIn(2.0, 4.5) && Cuts::pT > 20 * GeV &&
                                 Cuts::abspid == PID::MUON, Cuts::massIn(60*GeV, 120*GeV));
      declare(zmumufinder, "ZmumuFinder");

      // Book histograms
      // specify custom binning
      book(_h_sigma_vs_y,    18, 1, 1);
      book(_h_sigma_vs_pt,   19, 1, 1);
      book(_h_sigma_vs_phi,  20, 1, 1);
      book(_h_sigma_vs_ypt,  {2., 2.5, 3., 3.5, 4., 4.5});
      book(_h_sigma_vs_yphi, {2., 2.5, 3., 3.5, 4., 4.5});
      for (size_t i = 1; i < _h_sigma_vs_ypt->numBins()+1; ++i) {
      	book(_h_sigma_vs_ypt->bin(i),  21, 1, i);
      	book(_h_sigma_vs_yphi->bin(i), 22, 1, i);
      }
    }


    /// Perform the per-event analysis
    void analyze(const Event& event) {

      // Retrieve dressed leptons, sorted by pT
      const DileptonFinder& zmumufinder = apply<DileptonFinder>(event, "ZmumuFinder");
      if(zmumufinder.empty()) vetoEvent;
      if(zmumufinder.bosons().size() > 1)
        MSG_WARNING("Found multiple (" << zmumufinder.bosons().size() << ") Z -> mu+ mu- decays!");

      // Z momenta
      FourMomentum zmumu = zmumufinder.bosons()[0].momentum();
      if (zmumufinder.leptons().size() < 2) vetoEvent;

      const Particle& muon_p = zmumufinder.constituents()[0];
      const Particle& muon_m = zmumufinder.constituents()[1];

      const double diffphi = deltaPhi(muon_p, muon_m);
      const double diffpsd = deltaEta(muon_p, muon_m);
      const double accphi = M_PI - diffphi;
      const double angular = tan(accphi/2) / cosh(diffpsd/2);

      _h_sigma_vs_y->fill(zmumu.absrapidity());
      _h_sigma_vs_pt->fill(zmumu.pT()/GeV);
      _h_sigma_vs_phi->fill(angular);
      _h_sigma_vs_ypt->fill(zmumu.absrapidity(), zmumu.pT()/GeV);
      _h_sigma_vs_yphi->fill(zmumu.absrapidity(), angular);
    }


    /// Normalise histograms etc., after the run
    void finalize() {
      const double xs = crossSection()/picobarn;
      double scale_f = xs/sumOfWeights()/2.;
      _h_sigma_vs_y->scaleW(scale_f);
      _h_sigma_vs_pt->scaleW(scale_f);
      _h_sigma_vs_phi->scaleW(scale_f);
      scale(_h_sigma_vs_ypt,  scale_f*2.);
      scale(_h_sigma_vs_yphi, scale_f*2.);
    }

    ///@}


    Histo1DPtr _h_sigma_vs_y, _h_sigma_vs_pt, _h_sigma_vs_phi;
    Histo1DGroupPtr _h_sigma_vs_ypt, _h_sigma_vs_yphi;

  };


  RIVET_DECLARE_PLUGIN(LHCB_2021_I1990313);

}
