#include "Rivet/Analysis.hh"
#include "Rivet/Projections/FastJets.hh"
#include "Rivet/Projections/FinalState.hh"
#include "Rivet/Projections/LeptonFinder.hh"
#include "Rivet/Projections/PromptFinalState.hh"
#include "Rivet/Projections/VetoedFinalState.hh"

namespace Rivet {


  /// ttbb at 13 TeV
  class ATLAS_2018_I1705857 : public Analysis {
  public:

    /// Constructor
    RIVET_DEFAULT_ANALYSIS_CTOR(ATLAS_2018_I1705857);


    void init() {
      // Eta ranges
      Cut eta_full = (Cuts::abseta < 5.0);
      // Lepton cuts
      Cut lep_cuts25 = (Cuts::abseta < 2.5) && (Cuts::pT >= 25*GeV);
      // All final state particles
      FinalState fs(eta_full);

      // Get photons to dress leptons
      PromptFinalState photons(eta_full && Cuts::abspid == PID::PHOTON, TauDecaysAs::PROMPT);

      // Projection to find the electrons
      PromptFinalState electrons(eta_full && Cuts::abspid == PID::ELECTRON, TauDecaysAs::PROMPT);

      // Projection to find the muons
      PromptFinalState muons(eta_full && Cuts::abspid == PID::MUON, TauDecaysAs::PROMPT);

      LeptonFinder dressedelectrons25(electrons, photons, 0.1, lep_cuts25);
      LeptonFinder dressedmuons25(muons, photons, 0.1, lep_cuts25);

      declare(dressedelectrons25, "elecs");
      declare(dressedmuons25, "muons");

      // From here on we are just setting up the jet clustering
      IdentifiedFinalState nu_id;
      nu_id.acceptNeutrinos();
      PromptFinalState neutrinos(nu_id);
      neutrinos.acceptTauDecays(true);

      PromptFinalState jet_photons(eta_full && Cuts::abspid == PID::PHOTON, TauDecaysAs::NONPROMPT);
      LeptonFinder all_dressed_electrons(electrons, jet_photons, 0.1, eta_full);
      LeptonFinder all_dressed_muons(muons, jet_photons, 0.1, eta_full);

      VetoedFinalState vfs(fs);
      vfs.addVetoOnThisFinalState(all_dressed_electrons);
      vfs.addVetoOnThisFinalState(all_dressed_muons);
      vfs.addVetoOnThisFinalState(neutrinos);

      FastJets jets(vfs, JetAlg::ANTIKT, 0.4, JetMuons::DECAY, JetInvisibles::DECAY);
      declare(jets, "jets");

      // fiducial cross-section histogram
      book(_histograms["fid_xsec"], 1, 1, 1);
      book(_histograms["fid_xsec_no_ttX"], 2, 1, 1);

      book(_d["nbjets_emu"], 3, 1, 1);
      book(_d["nbjets_emu_no_ttX"], 4, 1, 1);

      // HT
      book_hist("ht_emu", 3);
      book_hist("ht_had_emu", 4);

      book_hist("ht_ljets", 5);
      book_hist("ht_had_ljets", 6);

      // b-jet pTs
      book_hist("lead_bjet_pt_emu",    7);
      book_hist("sublead_bjet_pt_emu", 8);
      book_hist("third_bjet_pt_emu",   9);

      book_hist("lead_bjet_pt_ljets",    10);
      book_hist("sublead_bjet_pt_ljets", 11);
      book_hist("third_bjet_pt_ljets",   12);
      book_hist("fourth_bjet_pt_ljets",  13);

      // leading bb pair
      book_hist("m_bb_leading_emu",  14);
      book_hist("pt_bb_leading_emu", 15);
      book_hist("dR_bb_leading_emu", 16);

      book_hist("m_bb_leading_ljets",  17);
      book_hist("pt_bb_leading_ljets", 18);
      book_hist("dR_bb_leading_ljets", 19);

      // closest bb pair
      book_hist("m_bb_closest_emu",  20);
      book_hist("pt_bb_closest_emu", 21);
      book_hist("dR_bb_closest_emu", 22);

      book_hist("m_bb_closest_ljets",  23);
      book_hist("pt_bb_closest_ljets", 24);
      book_hist("dR_bb_closest_ljets", 25);
    }


    void analyze(const Event& event) {
      DressedLeptons leptons;
      for (auto &lep : apply<LeptonFinder>(event, "muons").dressedLeptons()) { leptons.push_back(lep); }
      for (auto &lep : apply<LeptonFinder>(event, "elecs").dressedLeptons()) { leptons.push_back(lep); }

      const Jets jets = apply<FastJets>(event, "jets").jetsByPt(Cuts::pT > 25*GeV && Cuts::abseta < 2.5);
      for (const auto& jet : jets) {
        idiscard(leptons, [&](const DressedLepton& lep) { return deltaR(jet, lep) < 0.4; });
      }

      Jets bjets;
      for (const Jet& jet : jets) {
        if (jet.bTagged(Cuts::pT >= 5*GeV))  bjets += jet;
      }

      size_t njets = jets.size();
      size_t nbjets = bjets.size();

      // Evaluate basic event selection
      bool pass_ljets = (leptons.size() == 1 && leptons[0].pT() > 27*GeV);

      bool pass_emu =
        // 2 leptons > 27 GeV
        (leptons.size() == 2) &&
        (leptons[0].pT() > 27*GeV && leptons[1].pT() > 27*GeV) &&
        // emu events
        ((leptons[0].abspid() == 11 && leptons[1].abspid() == 13) ||
         (leptons[0].abspid() == 13 && leptons[1].abspid() == 11)) &&
        // opposite charge
        (leptons[0].charge() != leptons[1].charge());

      // If we don't have exactly 1 or 2 leptons then veto the event
      if (!pass_emu && !pass_ljets)  vetoEvent;

      if (pass_emu) {
        if (nbjets >= 2)  dfill("nbjets_emu", nbjets);
        if (nbjets >= 3)  fill("fid_xsec", 1);
        if (nbjets >= 4)  fill("fid_xsec", 2);
      }

      if (pass_ljets) {
        if (nbjets >= 3 && njets >= 5)  fill("fid_xsec", 3);
        if (nbjets >= 4 && njets >= 6)  fill("fid_xsec", 4);
      }

      if (pass_emu && (nbjets < 3 || njets < 3))    vetoEvent;
      if (pass_ljets && (nbjets < 4 || njets < 6))  vetoEvent;

      double hthad = sum(jets, Kin::pT, 0.0);
      double ht = sum(leptons, Kin::pT, hthad);

      FourMomentum jsum = bjets[0].momentum() + bjets[1].momentum();
      double dr_leading = deltaR(bjets[0], bjets[1]);

      size_t ind1 = 0, ind2 = 0; double mindr = 999.;
      for (size_t i = 0; i < bjets.size(); ++i) {
        for (size_t j = 0; j < bjets.size(); ++j) {
          if (i == j)  continue;
          double dr = deltaR(bjets[i], bjets[j]);
          if (dr < mindr) {
            ind1 = i;
            ind2 = j;
            mindr = dr;
          }
        }
      }

      FourMomentum bb_closest = bjets[ind1].momentum() + bjets[ind2].momentum();
      double dr_closest = deltaR(bjets[ind1], bjets[ind2]);

      if (pass_ljets) {
        // b-jet pTs
        fill("lead_bjet_pt_ljets", bjets[0].pT()/GeV);
        fill("sublead_bjet_pt_ljets", bjets[1].pT()/GeV);
        fill("third_bjet_pt_ljets", bjets[2].pT()/GeV);

        if (nbjets >= 4)  fill("fourth_bjet_pt_ljets", bjets[3].pT()/GeV);

        // HT
        fill("ht_ljets", ht/GeV);
        fill("ht_had_ljets", hthad/GeV);

        // leading bb pair
        fill("m_bb_leading_ljets", jsum.mass()/GeV);
        fill("pt_bb_leading_ljets", jsum.pT()/GeV);
        fill("dR_bb_leading_ljets", dr_leading);

        // closest bb pair
        fill("m_bb_closest_ljets", bb_closest.mass()/GeV);
        fill("pt_bb_closest_ljets", bb_closest.pT()/GeV);
        fill("dR_bb_closest_ljets", dr_closest);
      }
      if (pass_emu) {
        // b-jet pTs
        fill("lead_bjet_pt_emu", bjets[0].pT()/GeV);
        fill("sublead_bjet_pt_emu", bjets[1].pT()/GeV);
        fill("third_bjet_pt_emu", bjets[2].pT()/GeV);

        // HT
        fill("ht_emu", ht/GeV);
        fill("ht_had_emu", hthad/GeV);

        // leading bb pair
        fill("m_bb_leading_emu", jsum.mass()/GeV);
        fill("pt_bb_leading_emu", jsum.pT()/GeV);
        fill("dR_bb_leading_emu", dr_leading);

        // closest bb pair
        fill("m_bb_closest_emu", bb_closest.mass()/GeV);
        fill("pt_bb_closest_emu", bb_closest.pT()/GeV);
        fill("dR_bb_closest_emu", dr_closest);
      }
    }

    void finalize() {
      // Normalise all histograms
      const double sf = crossSection() / femtobarn / sumOfWeights();
      for (auto& h : _histograms) {
        scale(h.second, sf);
        if (h.first.find("fid_xsec") != string::npos)  continue;
        normalize(h.second, 1.0);
      }
      scale(_d, sf);
      normalize(_d["nbjets_emu"]);
      normalize(_d["nbjets_emu_no_ttX"]);
    }

    void fill(const string& name, const double value) {
      _histograms[name]->fill(value);
      _histograms[name + "_no_ttX"]->fill(value);
    }

    void dfill(const string& name, const size_t value) {
      string edge("OTHER");
      if (value == 2)  edge = "2";
      else if (value == 3)  edge = "3";
      else if (value >= 4)  edge = ">= 4";
      _d[name]->fill(edge);
      _d[name + "_no_ttX"]->fill(edge);
    }

    void book_hist(const std::string& name, unsigned int d) {
      // ttX substracted histograms are even numbers
      book(_histograms[name], (d * 2) - 1, 1, 1);
      book(_histograms[name + "_no_ttX"], d * 2, 1, 1);
    }

    private:
      map<std::string, Histo1DPtr> _histograms;
      map<std::string, BinnedHistoPtr<string>> _d;

  };

  RIVET_DECLARE_PLUGIN(ATLAS_2018_I1705857);
}
