#include "CxAODTools_HH_bbtautau/HHbbtautauLepHadSelection.h"
#include "xAODTau/TauJet.h"
#include "xAODEgamma/Electron.h"
#include "xAODEgamma/Photon.h"
#include "xAODMuon/Muon.h"
#include "xAODJet/Jet.h"
#include "xAODEventInfo/EventInfo.h"

#include<TRandom.h>

void HHbbtautauLepHadSelection::clearResult() {
  m_result.pass = false;
  m_result.taus.clear();
  m_result.signalJets.clear();
  m_result.forwardJets.clear();
  m_result.met = nullptr;
  m_result.mu=nullptr; 
  m_result.el=nullptr;

}

bool HHbbtautauLepHadSelection::passSelection(SelectionContainers & containers,
					      bool isKinVar) {
  
  // here just as an example:
  // if a new passKinematics() function is defined with some variables in the prototype,
  // one needs to reimplement passSelection here
  // otherwise, don't need to put any code
 
  return HHbbtautauSelection<ResultHHbbtautauLepHad>::passSelection(containers, isKinVar);
}

bool HHbbtautauLepHadSelection::passTauSelection(const xAOD::TauJetContainer* taus,  const xAOD::EventInfo* evtinfo)
{
  int res=doHHTauSelection(taus);

  // Only do if antitau config (not separate slector as also needs to be done for DLV reversed and antiisol)
  res=doRandomTauSelection(evtinfo);

  if(res>=1) {
    return true;
  }
  return false;

}

bool HHbbtautauLepHadSelection::passLeptonSelection(const xAOD::ElectronContainer* electrons,
                                                    const xAOD::MuonContainer* muons,
                                                    const xAOD::MissingET* met)
{
  const xAOD::Electron* el_junk=nullptr;
  const xAOD::Muon* mu_junk=nullptr;
  int res = doHHLeptonSelection(electrons, muons, m_result.el, el_junk, m_result.mu, mu_junk);
  if(res != 1)//(res<2) //no DLV
    {
      return false;
    }
  m_result.met = met;
  return true;

}

// preselection versions
bool HHbbtautauLepHadSelection::passPreSelection(SelectionContainers & containers,
						 bool isKinVar) {
  
  // here just as an example:
  // if a new passKinematics() function is defined with some variables in the prototype,
  // one needs to reimplement passSelection here
  // otherwise, don't need to put any code
  return HHbbtautauSelection<ResultHHbbtautauLepHad>::passPreSelection(containers, isKinVar);
}

bool HHbbtautauLepHadSelection::passLeptonPreSelection(const xAOD::ElectronContainer* electrons,
                                               const xAOD::MuonContainer* muons,
                                               const xAOD::MissingET* met) {
  const xAOD::Electron* el_junk=nullptr;
  const xAOD::Muon* mu_junk=nullptr;
  int res = doHHLeptonPreSelection(electrons, muons, m_result.el, el_junk, m_result.mu, mu_junk);
  if(res != 1)//(res!=2) //inverse DLV. DLV is performed for loose leptons as well.
  {
     return false;
  }
  return true;
  
}

bool HHbbtautauLepHadSelection::passKinematics() {
  // MJ cuts, like MET / MPT etc...
  // my advice is to add in passKinematics() prototype all the stuff that
  // doesn't need to be put in the Result struct, like MPT

  return true;
}

EL::StatusCode HHbbtautauLepHadSelection::writeEventVariables(const xAOD::EventInfo* eventInfoIn,
                             xAOD::EventInfo* eventInfoOut,
                             bool isKinVar,
                             bool isWeightVar,
			     std::string sysName,
			     int rdm_RunNumber,
			     CP::MuonTriggerScaleFactors* trig_sfmuon) {

  // ensure to write all variables for all events
  Props::leptonSF.set(eventInfoOut, 1);
  
  if (!m_result.pass) return EL::StatusCode::SUCCESS;
  
  // lepton SF
  if (m_result.isEl()) {
    float leptonSF = 1;
    leptonSF *= Props::effSFmediumLH.get(m_result.el);
    Props::leptonSF.set(eventInfoOut, leptonSF);
  } else {
    float leptonSF = 1;
    leptonSF *= Props::effSF.get(m_result.mu);
    Props::leptonSF.set(eventInfoOut, leptonSF);
  }

  return EL::StatusCode::SUCCESS;
}

int HHbbtautauLepHadSelection::doRandomTauSelection(const xAOD::EventInfo* evtinfo) {

  //if (!m_antitau) return m_result.taus.size();
  
  bool passSLe(false), passSLm(false);
  // Make sure works when reverse DLV

  const xAOD::Electron* e = m_result.el;
  if (e && e->pt() > 25000. && 
      Props::passHLT_e24_lhmedium_iloose_L1EM20VH.get(evtinfo) &&
      Props::matchHLT_e24_lhmedium_iloose_L1EM20VH.get(e)) {
    passSLe = true;    
  }

  const xAOD::Muon* m = m_result.mu;
  if (m && m->pt() > 21000. && 
      Props::passHLT_mu20_iloose_L1MU15.get(evtinfo) &&
      Props::matchHLT_mu20_iloose_L1MU15.get(m)) {
    passSLm = true;
  }

  std::vector<const xAOD::TauJet*> taus = m_result.taus;
  // Loop over selected taus to find valid taus matched to a lep+tau trigger if not 
  // passed single-lep trigger; else all are valid

  int ntauSig(0);
  if (!passSLe && !passSLm) {
    for (auto itr = taus.begin(); itr != taus.end(); ) {
      const xAOD::TauJet* tau = *itr;
      if (passSLe && !(tau->pt() > 26000 &&
		       Props::passHLT_e17_medium_tau25_medium1_tracktwo.get(evtinfo) && 
		       Props::matchHLT_e17_medium_tau25_medium1_tracktwo.get(tau))) {
	Props::isHHSignalTau.set(tau, false);
	Props::forMETRebuild.set(tau, false);//Using this for MET rebuilding
	itr = taus.erase(itr);
      } else if (passSLm && !(tau->pt() > 36000 &&
			      Props::passHLT_mu14_tau35_medium1_tracktwo.get(evtinfo) && 
			      Props::matchHLT_mu14_tau35_medium1_tracktwo.get(tau))) {
	Props::isHHSignalTau.set(tau, false);
	Props::forMETRebuild.set(tau, false);//Using this for MET rebuilding
	itr = taus.erase(itr);
      } else {
	++itr;
	ntauSig++;
      }
    }
  } else {
    ntauSig = taus.size();
  }
  
  // Pick a random one from those remaining if more than one; remove the others

  if (ntauSig > 1) {
    int irand = gRandom->Integer(ntauSig);
    
    int itau(0);
    for (auto itr = taus.begin(); itr != taus.end(); ) {
      const xAOD::TauJet* tau = *itr;
      
      if (itau++ != irand) {
	Props::isHHSignalTau.set(tau, false);
	Props::forMETRebuild.set(tau, false);
	itr = taus.erase(itr);
      } else {
	++itr;
      }
    }
  }

  return 1;  // 1 selected random signal anti-tau by definition
}