#include "Rivet/Analysis.hh"
#include "Rivet/Projections/FinalState.hh"
#include "Rivet/Projections/FastJets.hh"
#include "Rivet/Projections/PartonicTops.hh"
#include "Rivet/Projections/LeptonFinder.hh"
#include "Rivet/Projections/IdentifiedFinalState.hh"
#include "Rivet/Projections/PromptFinalState.hh"
#include "Rivet/Projections/VetoedFinalState.hh"
#include "Rivet/Projections/InvMassFinalState.hh"
#include "Rivet/Projections/MissingMomentum.hh"

namespace Rivet {


  class CMS_2016_I1473674 : public Analysis {
  public:

    // Minimal constructor
    RIVET_DEFAULT_ANALYSIS_CTOR(CMS_2016_I1473674);


    // Set up projections and book histograms
    void init() {

      // Complete final state
      FinalState fs;

      // Parton level top quarks
      declare(PartonicTops(TopDecay::E_MU, PromptEMuFromTau::NO), "LeptonicPartonTops");
      declare(PartonicTops(TopDecay::HADRONIC), "HadronicPartonTops");

      // Projections for dressed electrons and muons
      IdentifiedFinalState photons(fs);
      photons.acceptIdPair(PID::PHOTON);
      //
      IdentifiedFinalState el_id(fs);
      el_id.acceptIdPair(PID::ELECTRON);
      PromptFinalState electrons(el_id);
      declare(electrons, "Electrons");
      LeptonFinder dressed_electrons(electrons, photons, 0.1);
      declare(dressed_electrons, "DressedElectrons");
      //
      IdentifiedFinalState mu_id(fs);
      mu_id.acceptIdPair(PID::MUON);
      PromptFinalState muons(mu_id);
      declare(muons, "Muons");
      LeptonFinder dressed_muons(muons, photons, 0.1);
      declare(dressed_muons, "DressedMuons");

      // Projection for jets
      VetoedFinalState fs_jets;
      fs_jets.addVetoOnThisFinalState(dressed_muons);
      declare(FastJets(fs_jets, JetAlg::ANTIKT, 0.5), "Jets");

      // Projections for MET
      declare(MissingMomentum(), "MET");


      // Booking of histograms
      book(_hist_met ,5, 1, 1);
      book(_hist_ht  ,6, 1, 1);
      book(_hist_st  ,7, 1, 1);
      book(_hist_wpt ,8, 1, 1);
    }


    /// Per-event analysis
    void analyze(const Event& event) {

      // Select ttbar -> lepton+jets at parton level, removing tau decays
      const Particles leptonicpartontops = apply<ParticleFinder>(event, "LeptonicPartonTops").particlesByPt();
      if (leptonicpartontops.size() != 1) vetoEvent;
      const Particles hadronicpartontops = apply<ParticleFinder>(event, "HadronicPartonTops").particlesByPt();
      if (hadronicpartontops.size() != 1) vetoEvent;

      // Select ttbar -> lepton+jets at particle level
      const LeptonFinder& dressed_electrons = apply<LeptonFinder>(event, "DressedElectrons");
      const LeptonFinder& dressed_muons = apply<LeptonFinder>(event, "DressedMuons");
      if (dressed_electrons.dressedLeptons().size() + dressed_muons.dressedLeptons().size() != 1) vetoEvent;
      const FourMomentum lepton = (dressed_electrons.dressedLeptons().empty() ? dressed_muons : dressed_electrons).dressedLeptons()[0];

      // MET
      const MissingMomentum& met = apply<MissingMomentum>(event, "MET");
      _hist_met->fill(met.visibleMomentum().pT()/GeV);

      // HT and ST
      const FastJets& jetpro = apply<FastJets>(event, "Jets");
      const Jets jets = jetpro.jetsByPt(Cuts::pT > 20*GeV);

      double ht = 0.0;
      for (const Jet& j : jets) {
        if (deltaR(j.momentum(), lepton) > 0.3) {
          ht += j.pT();
        }
      }

      double st = ht + lepton.pT() + met.visibleMomentum().pT();
      _hist_ht->fill(ht/GeV);
      _hist_st->fill(st/GeV);

      // WPT
      const FourMomentum w = lepton - met.visibleMomentum();
      _hist_wpt->fill(w.pT()/GeV);
    }


    /// Normalize histograms
    void finalize() {
      normalize(_hist_met);
      normalize(_hist_ht);
      normalize(_hist_st);
      normalize(_hist_wpt);
    }

  private:
    Histo1DPtr _hist_met, _hist_ht, _hist_st, _hist_wpt;

  };


  RIVET_DECLARE_PLUGIN(CMS_2016_I1473674);

}
