#include "SimpleAnalysisFramework/AnalysisClass.h"

DefineAnalysis(EwkTwoLeptonZeroJet2018)

void EwkTwoLeptonZeroJet2018::Init()
{
	addRegions({"DF_0J_a","DF_0J_b","DF_0J_c","DF_0J_d","DF_0J_e","DF_0J_f","DF_0J_g","DF_0J_h","DF_0J_i"});
	addRegions({"SF_0J_a","SF_0J_b","SF_0J_c","SF_0J_d","SF_0J_e","SF_0J_f","SF_0J_g","SF_0J_h","SF_0J_i"});
	addRegions({"DF_1J_a","DF_1J_b","DF_1J_c","DF_1J_d","DF_1J_e","DF_1J_f","DF_1J_g","DF_1J_h","DF_1J_i"});
	addRegions({"SF_1J_a","SF_1J_b","SF_1J_c","SF_1J_d","SF_1J_e","SF_1J_f","SF_1J_g","SF_1J_h","SF_1J_i"});
	addRegions({"DF_0J_Incl_a","DF_0J_Incl_b","DF_0J_Incl_c","DF_0J_Incl_d"});
	addRegions({"SF_0J_Incl_a","SF_0J_Incl_b","SF_0J_Incl_c","SF_0J_Incl_d"});
	addRegions({"DF_1J_Incl_a","DF_1J_Incl_b","DF_1J_Incl_c","DF_1J_Incl_d"});
	addRegions({"SF_1J_Incl_a","SF_1J_Incl_b","SF_1J_Incl_c","SF_1J_Incl_d"});
	addRegions({"CR_VZ","CR_WW","CR_top"});
}

void EwkTwoLeptonZeroJet2018::ProcessEvent(AnalysisEvent *event)
{
	auto baselineElectrons  = event->getElectrons(10, 2.47, ELooseBLLH | EZ05mm);
	auto baselineMuons       = event->getMuons(10, 2.6, MuMedium | MuNotCosmic | MuZ05mm | MuQoPSignificance); // 2L0J uses SUSY2 which skims on eta<2.6
	auto baselineJets       = event->getJets(20., 4.5);
	auto metVec             = event->getMET();
	double met              = metVec.Et();
	double metSignificance =  event->getMETSignificance();
	double eventWeight = event->getMCWeights()[0];

	// default OR
	auto radiusCalcLep = [] (const AnalysisObject& lep,const AnalysisObject&) {
							 return (0.04 + 10/lep.Pt()) > 0.4 ? 0.4 : 0.04 + 10/lep.Pt();
						 };

	// // ele-mu OR
	// baselineElectrons = overlapRemoval(baselineElectrons, baselineMuons, 0.01);
	baselineMuons = overlapRemoval(baselineMuons,baselineElectrons,0.01);
	// // ele-jet OR
	baselineJets = overlapRemoval(baselineJets, baselineElectrons, 0.2);
	baselineElectrons = overlapRemoval(baselineElectrons, baselineJets, radiusCalcLep);
	// mu-jet OR
	baselineJets = overlapRemoval(baselineJets, baselineMuons, 0.2, LessThan3Tracks);
	baselineMuons = overlapRemoval(baselineMuons, baselineJets, radiusCalcLep);

	// // // overlap removal described in 2L0J paper
	// baselineJets = overlapRemoval(baselineJets, baselineElectrons, 0.2);
	// baselineJets = overlapRemoval(baselineJets, baselineMuons, 0.4, LessThan3Tracks);
	// baselineElectrons = overlapRemoval(baselineElectrons, baselineJets, 0.4);
	// baselineMuons      = overlapRemoval(baselineMuons, baselineJets, 0.4);

	auto baselineLeptons = baselineElectrons + baselineMuons;
	int nBaselineLeptons = baselineLeptons.size();

	// Get signal electrons with FCLoose for pT < 200 GeV and FCHighPtCaloOnly for pT > 200
	// filterObjects() only allows lower pT cut, so work around that
	auto signalElectrons = filterObjects(baselineElectrons,25., 2.47, ETightLH | ED0Sigma5  | EZ05mm  | EIsoFCLoose);
	AnalysisObjects signalEleLowPt;
    for (const auto& ele : signalElectrons) {
      if ((ele.Pt() < 200.)) signalEleLowPt.push_back(ele);
  	}
	auto signalEleHighPt = filterObjects(baselineElectrons,200., 2.47, ETightLH | ED0Sigma5  | EZ05mm  | EIsoFCHighPtCaloOnly);
	signalElectrons = signalEleLowPt + signalEleHighPt;

	auto signalMuons = filterObjects(baselineMuons, 25, 2.6, MuD0Sigma3 | MuZ05mm | MuIsoFCLoose);

	auto signalJets = filterObjects(baselineJets, 20, 2.4, JVT120Jet);
	auto bJets = filterObjects(signalJets, 20., 2.4, BTag77MV2c10);
	auto nonbJets = filterObjects(signalJets, 20., 2.4, NOT(BTag77MV2c10));

	auto signalLeptons   = signalElectrons + signalMuons;

	int nJets = signalJets.size();
	int nBjets = bJets.size();
	int nNonBjets = nonbJets.size();
	int nLeptons = signalLeptons.size();
	//	int nElectrons = signalElectrons.size();
	//	int nMuons = signalMuons.size();

	sortObjectsByPt(signalLeptons);
	sortObjectsByPt(nonbJets);
	sortObjectsByPt(signalJets);
	sortObjectsByPt(signalElectrons);
	sortObjectsByPt(signalMuons);

	//start preselecting
	//exactly 2 baseline and signal leptons
	if(nBaselineLeptons != 2) return;
	if(nLeptons != 2) return;

	//lepton pTs > 25 GeV
	if(signalLeptons[0].Pt() < 25. || signalLeptons[1].Pt() < 25.) return;

	//require opposite charge leptons
	if(signalLeptons[0].charge() == signalLeptons[1].charge()) return;

	//computing relevant variables
	float mt2 = calcMT2(signalLeptons[0],signalLeptons[1],metVec);
	float mll = (signalLeptons[0] + signalLeptons[1]).M();
	bool isSF = true;
	if(signalLeptons[0].type() != signalLeptons[1].type()) isSF = false;

	// for debugging only
	double L2isEMU        = 0;
    double L2isEE         = 0;
    double L2isMUMU       = 0;

	if (signalLeptons[0].type() == AnalysisObjectType::ELECTRON && signalLeptons[1].type() == AnalysisObjectType::ELECTRON) L2isEE = 1;
    else if (signalLeptons[0].type() == AnalysisObjectType::MUON && signalLeptons[1].type() == AnalysisObjectType::MUON) L2isMUMU = 1;
    else L2isEMU = 1;

	// for debugging only
	float l1flav, l2flav;
	if ( signalLeptons[0].type()== AnalysisObjectType::ELECTRON ) l1flav = 11.0;
	else l1flav = 13.0;
	if ( signalLeptons[1].type()== AnalysisObjectType::ELECTRON ) l2flav = 11.0;
	else l2flav = 13.0;

	// debug ntup vars
	ntupVar("eventWeight",eventWeight);
	ntupVar("MET",met*1000);
	ntupVar("L2Mll",mll*1000);
	ntupVar("L2MT2",mt2*1000);
	ntupVar("L2isEMU",L2isEMU);
	ntupVar("L2isEE",L2isEE);
	ntupVar("L2isMUMU",L2isMUMU);
	ntupVar("nJet",nJets);
	ntupVar("L2nCentralBJet",nBjets);
	ntupVar("L2nCentralLightJet",nNonBjets);
	ntupVar("L2nForwardJet",0);
	ntupVar("nLeps",nLeptons);
	ntupVar("METsig",metSignificance);
	ntupVar("lept1Flav",l1flav);
	ntupVar("lept2Flav",l2flav);
	ntupVar("lept1Pt",signalLeptons[0].Pt()*1000);
	ntupVar("lept1Eta",signalLeptons[0].Eta());
	ntupVar("lept1Phi",signalLeptons[0].Phi());
	ntupVar("lept2Pt",signalLeptons[1].Pt()*1000);
	ntupVar("lept2Eta",signalLeptons[1].Eta());
	ntupVar("lept2Phi",signalLeptons[1].Phi());

	if (nNonBjets>0) {
		ntupVar("jet1Pt",signalJets[0].Pt()*1000);
		ntupVar("jet1Eta",signalJets[0].Eta());
		ntupVar("jet1Phi",signalJets[0].Phi());
	}

	//
	// Control regions
	//
	if (isSF && nNonBjets==0 && nBjets==0 && mt2>=120. && met>110. && metSignificance>10. && mll>=61.2 && mll<121.2) accept("CR_VZ");
	if (!isSF && nNonBjets==0 && nBjets==0 && mt2>=60. && mt2<65. && mll>=100. && met>=60. && met<100. && metSignificance>=5. && metSignificance<10.) accept("CR_WW");
	if (!isSF && nNonBjets==0 && nBjets==1 && mt2>=80. && met>110. && metSignificance>10. && mll>=100.) accept("CR_top");

	//vetoing on jets
	if(nBjets>0) return;
	if(nNonBjets>1) return;

	//MET preselection
	if(met<110.) return;

	//MET significance preselection
	if(metSignificance<10.) return;

	//all signal regions have mll > 100 GeV (DF) or mll > 121.2 GeV (SF)
	if (isSF && mll<121.2) return;
	else if (!isSF && mll<100.) return;

	if (nNonBjets == 0) {

		if(mt2 >= 100. && mt2 < 105. && !isSF) accept("DF_0J_a");
		else if(mt2 >= 105. && mt2 < 110. && !isSF) accept("DF_0J_b");
		else if(mt2 >= 110. && mt2 < 120. && !isSF) accept("DF_0J_c");
		else if(mt2 >= 120. && mt2 < 140. && !isSF) accept("DF_0J_d");
		else if(mt2 >= 140. && mt2 < 160. && !isSF) accept("DF_0J_e");
		else if(mt2 >= 160. && mt2 < 180. && !isSF) accept("DF_0J_f");
		else if(mt2 >= 180. && mt2 < 220. && !isSF) accept("DF_0J_g");
		else if(mt2 >= 220. && mt2 < 260. && !isSF) accept("DF_0J_h");
		else if(mt2 >= 260. && !isSF) accept("DF_0J_i");

		if(mt2 >= 100. && !isSF) accept("DF_0J_Incl_a");
		if(mt2 >= 160. && !isSF) accept("DF_0J_Incl_b");
		if(mt2 >= 100. && mt2 < 120. && !isSF) accept("DF_0J_Incl_c");
		if(mt2 >= 120. && mt2 < 160. && !isSF) accept("DF_0J_Incl_d");

		if(mt2 >= 100. && mt2 < 105. && isSF) accept("SF_0J_a");
		else if(mt2 >= 105. && mt2 < 110. && isSF) accept("SF_0J_b");
		else if(mt2 >= 110. && mt2 < 120. && isSF) accept("SF_0J_c");
		else if(mt2 >= 120. && mt2 < 140. && isSF) accept("SF_0J_d");
		else if(mt2 >= 140. && mt2 < 160. && isSF) accept("SF_0J_e");
		else if(mt2 >= 160. && mt2 < 180. && isSF) accept("SF_0J_f");
		else if(mt2 >= 180. && mt2 < 220. && isSF) accept("SF_0J_g");
		else if(mt2 >= 220. && mt2 < 260. && isSF) accept("SF_0J_h");
		else if(mt2 >= 260. && isSF) accept("SF_0J_i");

		if(mt2 >= 100. && isSF) accept("SF_0J_Incl_a");
		if(mt2 >= 160. && isSF) accept("SF_0J_Incl_b");
		if(mt2 >= 100. && mt2 < 120. && isSF) accept("SF_0J_Incl_c");
		if(mt2 >= 120. && mt2 < 160. && isSF) accept("SF_0J_Incl_d");

	}
	else {

		if(mt2 >= 100. && mt2 < 105. && !isSF) accept("DF_1J_a");
		if(mt2 >= 105. && mt2 < 110. && !isSF) accept("DF_1J_b");
		if(mt2 >= 110. && mt2 < 120. && !isSF) accept("DF_1J_c");
		if(mt2 >= 120. && mt2 < 140. && !isSF) accept("DF_1J_d");
		if(mt2 >= 140. && mt2 < 160. && !isSF) accept("DF_1J_e");
		if(mt2 >= 160. && mt2 < 180. && !isSF) accept("DF_1J_f");
		if(mt2 >= 180. && mt2 < 220. && !isSF) accept("DF_1J_g");
		if(mt2 >= 220. && mt2 < 260. && !isSF) accept("DF_1J_h");
		if(mt2 >= 260. && !isSF) accept("DF_1J_i");

		if(mt2 >= 100. && !isSF) accept("DF_1J_Incl_a");
		if(mt2 >= 160. && !isSF) accept("DF_1J_Incl_b");
		if(mt2 >= 100. && mt2 < 120. && !isSF) accept("DF_1J_Incl_c");
		if(mt2 >= 120. && mt2 < 160. && !isSF) accept("DF_1J_Incl_d");

		if(mt2 >= 100. && mt2 < 105. && isSF) accept("SF_1J_a");
		if(mt2 >= 105. && mt2 < 110. && isSF) accept("SF_1J_b");
		if(mt2 >= 110. && mt2 < 120. && isSF) accept("SF_1J_c");
		if(mt2 >= 120. && mt2 < 140. && isSF) accept("SF_1J_d");
		if(mt2 >= 140. && mt2 < 160. && isSF) accept("SF_1J_e");
		if(mt2 >= 160. && mt2 < 180. && isSF) accept("SF_1J_f");
		if(mt2 >= 180. && mt2 < 220. && isSF) accept("SF_1J_g");
		if(mt2 >= 220. && mt2 < 260. && isSF) accept("SF_1J_h");
		if(mt2 >= 260. && isSF) accept("SF_1J_i");

		if(mt2 >= 100. && isSF) accept("SF_1J_Incl_a");
		if(mt2 >= 160. && isSF) accept("SF_1J_Incl_b");
		if(mt2 >= 100. && mt2 < 120. && isSF) accept("SF_1J_Incl_c");
		if(mt2 >= 120. && mt2 < 160. && isSF) accept("SF_1J_Incl_d");

	}

	return;
}
