// -*- C++ -*-
//
// Package:    ExoDiBosonResonances/EDBRGenStudies
// Class:      EDBRGenStudies
// 
/**\class EDBRGenStudies EDBRGenStudies.cc ExoDiBosonResonances/EDBRGenStudies/plugins/EDBRGenStudies.cc

 Description: [one line class summary]

 Implementation:
     [Notes on implementation]
*/
//
// Original Author:  Jose Cupertino Ruiz Vargas
//         Created:  Wed, 17 Sep 2014 09:20:13 GMT
//
//


// system include files
#include <memory>

// user include files
#include "FWCore/Framework/interface/Frameworkfwd.h"
#include "FWCore/Framework/interface/EDAnalyzer.h"

#include "FWCore/Framework/interface/Event.h"
#include "FWCore/Framework/interface/MakerMacros.h"

#include "FWCore/ParameterSet/interface/ParameterSet.h"

#include "FWCore/ServiceRegistry/interface/Service.h"
#include "CommonTools/UtilAlgos/interface/TFileService.h"

#include "DataFormats/PatCandidates/interface/Electron.h"
#include "DataFormats/HepMCCandidate/interface/GenParticle.h"
#include "DataFormats/Common/interface/Association.h"
#include "DataFormats/Math/interface/deltaR.h"

#include "TTree.h"
#include "TEfficiency.h"

//
// class declaration
//

class EDBRGenStudies : public edm::EDAnalyzer {
   public:
      explicit EDBRGenStudies(const edm::ParameterSet&);
      ~EDBRGenStudies();

   private:
      virtual void beginJob() override;
      virtual void analyze(const edm::Event&, const edm::EventSetup&) override;
      virtual void endJob() override;

      // ----------member data ---------------------------
      edm::EDGetTokenT<edm::View<pat::Electron> > electronCollectionToken_;
      edm::EDGetTokenT<edm::Association<std::vector<reco::GenParticle> > > associatedGenParticleToken_;
      edm::InputTag electronIdTag_;

      TTree* ElectronTree; 
      double pt;
      double eta;
      double eeDeltaR;
      double ggDeltaR;
      double relIsoWithDBeta;
      int mult;
      int passID;
      int matchID;
      int statusOfMother;
      int pdgIdOfMother;

      TEfficiency *eff_pt;
      TEfficiency *fak_pt;
      TEfficiency *eff_eta;
      TEfficiency *fak_eta;
      TEfficiency *eff_dR;
      TEfficiency *eff_dR_pt;
};

//
// constants, enums and typedefs
//
typedef edm::Association<std::vector<reco::GenParticle> > MCTruthMap;

enum matchingCategory { UNMATCHED = 0, 
                        TRUE_PROMPT_ELECTRON,
                        TRUE_ELECTRON_FROM_TAU, 
                        TRUE_ELECTRON_FROM_Z, 
                        UNDEFINED }; //Anything else

//
// constructors and destructor
//
EDBRGenStudies::EDBRGenStudies(const edm::ParameterSet& iConfig):
  electronCollectionToken_(consumes<edm::View<pat::Electron> >(iConfig.getParameter<edm::InputTag>("electrons"))),
  associatedGenParticleToken_(consumes<MCTruthMap>(iConfig.getParameter<edm::InputTag>("associatedGenParticle"))),
  electronIdTag_(iConfig.getParameter<edm::InputTag>("electronIDs"))
{
   //now do what ever initialization is needed
   edm::Service<TFileService> fs;
   ElectronTree = fs->make<TTree>("ElectronTree","MC Truth matching");
   ElectronTree->Branch("pt",              &pt,              "pt/D");
   ElectronTree->Branch("eta",             &eta,             "eta/D");
   ElectronTree->Branch("eeDeltaR",        &eeDeltaR,        "eeDeltaR/D");
   ElectronTree->Branch("ggDeltaR",        &ggDeltaR,        "ggDeltaR/D");
   ElectronTree->Branch("relIsoWithDBeta", &relIsoWithDBeta, "relIsoWithDBeta/D");
   ElectronTree->Branch("mult",            &mult,            "mult/I");
   ElectronTree->Branch("passID",          &passID,          "passID/I");
   ElectronTree->Branch("matchID",         &matchID,         "matchID/I");
   ElectronTree->Branch("statusOfMother",  &statusOfMother,  "statusOfMother/I");
   ElectronTree->Branch("pdgIdOfMother",   &pdgIdOfMother,   "pdgIdOfMother/I");
}

EDBRGenStudies::~EDBRGenStudies()
{
}


//
// member functions
//

// ------------ method called for each event  ------------
void
EDBRGenStudies::analyze(const edm::Event& iEvent, const edm::EventSetup& iSetup)
{
   using namespace edm;

   Handle<View<pat::Electron> > collection;
   iEvent.getByToken(electronCollectionToken_, collection);
   const auto& ele_refs = collection->refVector();
   mult=0;

   Handle<MCTruthMap> mcMatchH;
   iEvent.getByToken(associatedGenParticleToken_, mcMatchH);
   const MCTruthMap & mcMatch = *mcMatchH;

   //****************************************************************//
   //                  Begin loop over pat::Electrons                //
   //****************************************************************//                                             
   for( const auto& el : ele_refs ) {
      // Isolation with delta beta correction
      reco::GsfElectron::PflowIsolationVariables pfIso = el->pfIsolationVariables();
      double absiso = pfIso.sumChargedHadronPt + std::max(0.0, pfIso.sumNeutralHadronEt + pfIso.sumPhotonEt - 0.5*pfIso.sumPUPt );
      relIsoWithDBeta = absiso/el->pt();
      // MC Truth Matching
      reco::GenParticleRef elTruth = mcMatch[el];
      if ( elTruth.isNonnull() ) {
           pt = elTruth->pt();
           eta = elTruth->eta();
           if ( elTruth->numberOfMothers() ) {
                mult++;
                const reco::GenParticle *genMother = (reco::GenParticle*)elTruth->mother(0);
                statusOfMother = genMother->status();
                pdgIdOfMother  = genMother->pdgId();
                switch ( abs(pdgIdOfMother) ) {
                   case 11: if(genMother->mother(0)->pdgId()==23) {
                               matchID = TRUE_ELECTRON_FROM_Z; 
                               const reco::GenParticle *genZ = (reco::GenParticle*)genMother->mother(0);
                               const reco::GenParticle *firstElectron  = (reco::GenParticle*)genZ->daughter(0)->daughter(0);
                               const reco::GenParticle *secondElectron = (reco::GenParticle*)genZ->daughter(1)->daughter(0);
                               ggDeltaR = reco::deltaR(firstElectron->p4(),secondElectron->p4());
                            }
                            else {
                               matchID = UNDEFINED; 
                               ggDeltaR = -99;
                            }
                            break; 
                   case 15: matchID = TRUE_ELECTRON_FROM_TAU; 
                            ggDeltaR = -99;
                            break; 
                   case 23: matchID = TRUE_ELECTRON_FROM_Z;   
                            ggDeltaR = reco::deltaR(genMother->daughter(0)->p4(),genMother->daughter(1)->p4());
                            break; 
                   default: matchID = UNDEFINED;              
                            ggDeltaR = -99;
                            break;  
                }
           }
           else {
                matchID = TRUE_PROMPT_ELECTRON; 
                statusOfMother = -99; //Dummy value
                pdgIdOfMother  = -99; //Dummy value
                ggDeltaR       = -99; //Dummy value
           }
           // deltaR distance (genE,recoE)
           const reco::Candidate::LorentzVector & recEl = el->p4();
           const reco::Candidate::LorentzVector & genEl = elTruth->p4();
           eeDeltaR = reco::deltaR(recEl,genEl);
      } 
      else {
           matchID = UNMATCHED; 
           pt             = -99;
           eta            = -99;
           statusOfMother = -99;  //Dummy value
           pdgIdOfMother  = -99;  //Dummy value
           ggDeltaR       = -99.; //Dummy value
           eeDeltaR       = -99.; //Dummy value
      }
      // End MC Truth Matching
      // Electron ID 
      passID = ( el->electronID(electronIdTag_.encode())>0.5 );
      // Efficiency and fake rate
      if (matchID == TRUE_ELECTRON_FROM_Z){ 
         eff_pt->Fill(passID,pt); 
         eff_eta->Fill(passID,eta);
         eff_dR->Fill(passID,ggDeltaR);
         eff_dR_pt->Fill(passID,pt,ggDeltaR);
      }
      if (matchID != TRUE_ELECTRON_FROM_Z){ 
         fak_pt->Fill(passID,pt); 
         fak_eta->Fill(passID,eta);
      }
      // Fill tree
      ElectronTree->Fill();
   }
   //****************************************************************//
   //                  End loop over pat::Electrons                  //
   //****************************************************************//
}


// ------------ method called once each job just before starting event loop  ------------
void 
EDBRGenStudies::beginJob()
{
   edm::Service<TFileService> fs;
   int nbins=10;
   double xbins[11] = {0.,10.,20.,50.,100.,200.,300.,400.,500.,700.,1000.};
   eff_pt = fs->make<TEfficiency>("eff_pt",";electron p_{T} (GeV);efficiency",nbins,xbins);
   fak_pt = fs->make<TEfficiency>("fak_pt",";electron p_{T} (GeV);fake rate", nbins,xbins);
   eff_eta = fs->make<TEfficiency>("eff_eta",";#eta;efficiency", 15,-3.,3.);
   fak_eta = fs->make<TEfficiency>("fak_eta",";#eta;fake rate",  15,-3.,3.);
   eff_dR = fs->make<TEfficiency>("eff_dR",";deltaR(genEle1,genEle2);efficiency", 15,0.,1.5);
   eff_dR_pt = fs->make<TEfficiency>("eff_dR_pt",";electron p_{T};deltaR(genEle1,genEle2)",9, 0.,900.,6,0.,0.6);
}

// ------------ method called once each job just after ending the event loop  ------------
void 
EDBRGenStudies::endJob() 
{
}

//define this as a plug-in
DEFINE_FWK_MODULE(EDBRGenStudies);