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

namespace Rivet {


  /// @brief Measurements of differential Z boson production cross sections in proton-proton collisions at 13 TeV
  class CMS_2019_I1753680 : public Analysis {
  public:

    /// Constructor
    RIVET_DEFAULT_ANALYSIS_CTOR(CMS_2019_I1753680);


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

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

      // Get options from the new option system
      // default to combined.
      _mode = 2;
      if ( getOption("LMODE") == "EL" ) _mode = 0;
      if ( getOption("LMODE") == "MU" ) _mode = 1;
      if ( getOption("LMODE") == "EMU" ) _mode = 2;

      // Initialise and register projections
      Cut cut = Cuts::abseta < 2.4 && Cuts::pT > 25*GeV;
      DileptonFinder zeeFind(91.2*GeV, 0.1, cut && Cuts::abspid == PID::ELECTRON, Cuts::massIn(76.2*GeV, 106.2*GeV));
      declare(zeeFind, "ZeeFind");
      DileptonFinder zmmFind(91.2*GeV, 0.1, cut && Cuts::abspid == PID::MUON    , Cuts::massIn(76.2*GeV, 106.2*GeV));
      declare(zmmFind, "ZmmFind");

      // Book histograms
      book(_h_Zmm_absY,    26, 1, 1);
      book(_h_Zee_absY,    26, 1, 2);
      book(_h["absY"],     26, 1, 3);
      book(_h_Zmm_pt,      27, 1, 1);
      book(_h_Zee_pt,      27, 1, 2);
      book(_h["pt"],       27, 1, 3);
      book(_h_Zmm_phiStar, 28, 1, 1);
      book(_h_Zee_phiStar, 28, 1, 2);
      book(_h["phiStar"],  28, 1, 3);
      book(_h["pt_Y0"],    29, 1, 1);
      book(_h["pt_Y1"],    29, 1, 2);
      book(_h["pt_Y2"],    29, 1, 3);
      book(_h["pt_Y3"],    29, 1, 4);
      book(_h["pt_Y4"],    29, 1, 5);

      book(_h_norm["pt"],      30, 1, 1);
      book(_h_norm["phiStar"], 31, 1, 1);
      book(_h_norm["absY"],    32, 1, 1);
      book(_h_norm["pt_Y0"],   33, 1, 1);
      book(_h_norm["pt_Y1"],   33, 1, 2);
      book(_h_norm["pt_Y2"],   33, 1, 3);
      book(_h_norm["pt_Y3"],   33, 1, 4);
      book(_h_norm["pt_Y4"],   33, 1, 5);

    }


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

      const DileptonFinder& zeeFS = apply<DileptonFinder>(event, "ZeeFind");
      const DileptonFinder& zmumuFS = apply<DileptonFinder>(event, "ZmmFind");

      const Particles& zees = zeeFS.bosons();
      const Particles& zmumus = zmumuFS.bosons();

      if (zees.size() + zmumus.size() != 1) {
        MSG_DEBUG("Did not find exactly one good Z candidate");
        vetoEvent;
      }

      //event identification depending on mass window
      bool ee_event=false;
      bool mm_event=false;

      if (zees.size() == 1) {
        ee_event = true;
      }
      if (zmumus.size() == 1) {
        mm_event = true;
      }

      if (ee_event && _mode == 1)
        vetoEvent;
      if (mm_event && _mode == 0)
        vetoEvent;

      const Particles& theLeptons = ee_event ? zeeFS.constituents() : zmumuFS.constituents();
      const Particle& lminus = theLeptons[0].charge() < 0 ? theLeptons[0] : theLeptons[1];
      const Particle& lplus = theLeptons[0].charge() < 0 ? theLeptons[1] : theLeptons[0];

      //calculate phi*
      const double thetaStar = acos(tanh( 0.5 * (lminus.eta() - lplus.eta()) ));
      const double dPhi = M_PI - deltaPhi(lminus, lplus);
      const double phiStar = tan(0.5 * dPhi) * sin(thetaStar);

      const Particle& zcand = ee_event ? zees[0] : zmumus[0];

      if (ee_event) {
        _h_Zee_absY->fill(zcand.absrap());
        _h_Zee_pt->fill(zcand.pt()/GeV);
        _h_Zee_phiStar->fill(phiStar);
      }
      else if (mm_event) {
        _h_Zmm_absY->fill(zcand.absrap());
        _h_Zmm_pt->fill(zcand.pt()/GeV);
        _h_Zmm_phiStar->fill(phiStar);
      }

      _h["pt"]->fill(zcand.pt()/GeV);
      _h_norm["pt"]->fill(zcand.pt()/GeV);
      _h["phiStar"]->fill(phiStar);
      _h_norm["phiStar"]->fill(phiStar);
      _h["absY"]->fill(zcand.absrap());
      _h_norm["absY"]->fill(zcand.absrap());

      if (zcand.absrap() < 0.4) {
        _h["pt_Y0"]->fill(zcand.pt()/GeV);
        _h_norm["pt_Y0"]->fill(zcand.pt()/GeV);
      }
      else if (zcand.absrap() < 0.8) {
        _h["pt_Y1"]->fill(zcand.pt()/GeV);
        _h_norm["pt_Y1"]->fill(zcand.pt()/GeV);
      }
      else if (zcand.absrap() < 1.2) {
        _h["pt_Y2"]->fill(zcand.pt()/GeV);
        _h_norm["pt_Y2"]->fill(zcand.pt()/GeV);
      }
      else if (zcand.absrap() < 1.6) {
        _h["pt_Y3"]->fill(zcand.pt()/GeV);
        _h_norm["pt_Y3"]->fill(zcand.pt()/GeV);
      }
      else if (zcand.absrap() < 2.4) {
        _h["pt_Y4"]->fill(zcand.pt()/GeV);
        _h_norm["pt_Y4"]->fill(zcand.pt()/GeV);
      }

    }

    /// Normalise histograms etc., after the run
    void finalize() {

      double norm = (sumOfWeights() != 0) ? crossSection()/picobarn/sumOfWeights() : 1.0;

      scale(_h_Zmm_pt,      norm);
      scale(_h_Zmm_absY,    norm);
      scale(_h_Zmm_phiStar, norm);

      scale(_h_Zee_pt,      norm);
      scale(_h_Zee_absY,    norm);
      scale(_h_Zee_phiStar, norm);

      // when running in combined mode, need to average to get lepton xsec
      if (_mode == 2) norm /= 2.;
      scale(_h, norm);

      for (auto& item : _h_norm) {
        const double rho = item.second->densitySum(false);
        if (rho)  scale(item.second, 1.0/rho);
      }

    }

    /// @}

  protected:

    size_t _mode;

    /// @name Histograms

  private:

    Histo1DPtr   _h_Zmm_pt, _h_Zmm_phiStar, _h_Zmm_absY;
    Histo1DPtr   _h_Zee_pt, _h_Zee_phiStar, _h_Zee_absY;

    map<string,Histo1DPtr> _h;
    map<string,Histo1DPtr> _h_norm;

  };


  RIVET_DECLARE_PLUGIN(CMS_2019_I1753680);


}
