#include "CxAODMaker_HH_bbtautau/TauHandler_HH_bbtautau.h"
#include <iostream>
#include <TRandom.h>

#include "xAODEgamma/Electron.h"
#include "xAODEgamma/ElectronContainer.h"
#include "xAODEgamma/ElectronAuxContainer.h"
#include "xAODMuon/Muon.h"
#include "xAODMuon/MuonContainer.h"
#include "xAODMuon/MuonAuxContainer.h"


TauHandler_HH_bbtautau::TauHandler_HH_bbtautau(const std::string& name, ConfigStore & config, xAOD::TEvent *event,
               EventInfoHandler & eventInfoHandler) :
  TauHandler (name, config, event, eventInfoHandler)
{
  using std::placeholders::_1; 
  m_selectFcns.clear(); 
  m_selectFcns.push_back(std::bind(&TauHandler_HH_bbtautau::LooseTau, this, _1)); 
  m_selectFcns.push_back(std::bind(&TauHandler_HH_bbtautau::passHH_bbtautauTau, this, _1)); 
}
bool TauHandler_HH_bbtautau::LooseTau(xAOD::TauJet * tau) {


  bool passSel = true;
  if( tau->pt() < 20000. ) passSel=false;  
  if (passSel) m_cutflow->count("Pt Cut");
  if(  fabs( tau->eta()  ) > 2.46 ) passSel=false;
  if (passSel) m_cutflow->count("Eta Cut");
  if (!(Props::passTauSelector.get(tau))) passSel = false;
  if (passSel) m_cutflow->count("TauSelector");

  if(m_antitau){  
    if (Props::BDTScore.get(tau) > 0.35 && !(Props::isBDTMedium.get(tau))) passSel = false;
  }
  
  Props::passPreSel.set(tau,passSel);
  Props::isHHLooseTau.set(tau,passSel);
  Props::forMETRebuild.set(tau, passSel);//Using this for MET rebuilding
  return passSel;
}
bool TauHandler_HH_bbtautau::passHH_bbtautauTau(xAOD::TauJet * tau)
{
  // Same as Loose Taus
  bool passSel = true;
  if (!(Props::passPreSel.get(tau))) passSel = false;
  if (!(Props::isBDTMedium.get(tau))) passSel = false; //Signal Region (not in the config anymore so that we play around with CRs)
  Props::isHHSignalTau.set(tau, passSel);
 
  return passSel;
}
EL::StatusCode TauHandler_HH_bbtautau::setVariablesCustom(xAOD::TauJet* inTau, xAOD::TauJet* outTau, bool isKinVar, bool isWeightVar)
{
  // set something without having a pre-defined method
  Props::isHHLooseTau.copy(inTau, outTau);
  Props::isHHSignalTau.copy(inTau, outTau);
  
  return EL::StatusCode::SUCCESS;

}

EL::StatusCode TauHandler_HH_bbtautau::selectRandomTau(std::map<TString,xAOD::ElectronContainer*>* electrons, std::map<TString,xAOD::MuonContainer*>* muons) 
{
  
  if(! m_antitau) return EL::StatusCode::SUCCESS;

  if (m_msgLevel == MSG::DEBUG) {
    Info("TauHandler::selectRandomTau()", "Called");
  }
  
  // Loop over variations
  for (std::pair<TString,xAOD::TauJetContainer*> particles : m_inContainer) {
    const xAOD::ElectronContainer* es = (*electrons)[particles.first];
    const xAOD::MuonContainer* mus = (*muons)[particles.first];
    
    bool passSL(false);
    int neSig(0), nmSig(0);
    
    // Should we just use signal?
    // What about single tau triggers at high pt (for tau had)?
    // Make sure works when reverse DLV

    for (auto e : *es) {
      if (!Props::passOR.get(e)) continue;
      if (!Props::isHHLooseElectron.get(e)) continue;
      if (e->pt() > 25000. && 
	  Props::passHLT_e24_lhmedium_iloose_L1EM20VH.get(m_eventInfoHandler.getEventInfo()) &&
	  Props::matchHLT_e24_lhmedium_iloose_L1EM20VH.get(e)) {
	passSL = true;    
      }
      if (Props::isBBTTSignalElectron.get(e)) neSig++;
    }
    
    for (auto m : *mus) {
      if (!Props::passOR.get(m)) continue;
      if (!Props::isHHLooseMuon.get(m)) continue;
      if (m->pt() > 21000. && 
	  Props::passHLT_mu20_iloose_L1MU15.get(m_eventInfoHandler.getEventInfo()) &&
	  Props::matchHLT_mu20_iloose_L1MU15.get(m)) {
	passSL = true;
      }
      if (Props::isBBTTSignalMuon.get(m)) nmSig++;
    }

    // Would be better to really do in event selection and then do the MET after this else may change and get out of sync
    // Could have a flag to do MET after event selection    

    // These events will anyway get removed -> what default to in this case?
    if (neSig > 0  and nmSig > 0) continue;
    
    
    // Loop over selected taus to find taus matched to a lep+tau trigger if not passed single-lep trigger;
    // else all are valid
    int ntauSig(0);
    if (!passSL) {
      for (auto tau : *particles.second) {
	if (!Props::passOR.get(tau)) continue;
	if (!Props::isHHSignalTau.get(tau)) continue;
	
	if (neSig > 0 && !(tau->pt() > 26000 &&
			   Props::passHLT_e17_medium_tau25_medium1_tracktwo.get(m_eventInfoHandler.getEventInfo()) && 
			   Props::matchHLT_e17_medium_tau25_medium1_tracktwo.get(tau))) {
	  Props::isHHSignalTau.set(tau, false);
	  Props::forMETRebuild.set(tau, false);//Using this for MET rebuilding
	} else if (nmSig > 0 && !(tau->pt() > 36000 &&
				  Props::passHLT_mu14_tau35_medium1_tracktwo.get(m_eventInfoHandler.getEventInfo()) && 
				  Props::matchHLT_mu14_tau35_medium1_tracktwo.get(tau))) {
	  Props::isHHSignalTau.set(tau, false);
	  Props::forMETRebuild.set(tau, false);//Using this for MET rebuilding
	} else {
	  ntauSig++;
	}
      }
    } else {
      ntauSig = particles.second->size();
    }
    
    // Pick a random one from those remaining
    int irand = gRandom->Integer(ntauSig);
    
    int itau(0);
    for (auto tau : *particles.second) {
      if (itau++ != irand) {
	Props::isHHSignalTau.set(tau, false);
	Props::forMETRebuild.set(tau, false);
      }
    }
  }
  return EL::StatusCode::SUCCESS;
}