#include "SimpleAnalysisFramework/AnalysisClass.h"
#include <string>

DefineAnalysis(OneLepton2018)

// Glance: ANA-SUSY-2018-10,
// Support note: ATL-COM-PHYS-2018-1638

void OneLepton2018::Init()
{
    // Exclusion SRs
    addRegions({"SR2JBVbin1","SR2JBVbin2","SR2JBVbin3"});
    addRegions({"SR2JBTbin1","SR2JBTbin2","SR2JBTbin3"});
    addRegions({"SR4JhxBVbin1","SR4JhxBVbin2","SR4JhxBVbin3"});
    addRegions({"SR4JhxBTbin1","SR4JhxBTbin2","SR4JhxBTbin3"});
    addRegions({"SR4JlxBVbin1","SR4JlxBVbin2","SR4JlxBVbin3"});
    addRegions({"SR4JlxBTbin1","SR4JlxBTbin2","SR4JlxBTbin3"});
    addRegions({"SR6JBVbin1","SR6JBVbin2","SR6JBVbin3","SR6JBVbin4"});
    addRegions({"SR6JBTbin1","SR6JBTbin2","SR6JBTbin3","SR6JBTbin4"});

    // Discovery SRs
    addRegions({"SR2Jhigh","SR2Jlow"});
    addRegions({"SR4Jlx","SR4Jhx"});
    addRegions({"SR6Jhigh","SR6Jlow"});

    // CRs
    addRegions({"WR2Jbin1","WR2Jbin2","WR2Jbin3"});
    addRegions({"TR2Jbin1","TR2Jbin2","TR2Jbin3"});
    addRegions({"WR4Jbin1","WR4Jbin2","WR4Jbin3"});
    addRegions({"TR4Jbin1","TR4Jbin2","TR4Jbin3"});
    addRegions({"WR6Jbin1","WR6Jbin2","WR6Jbin3","WR6Jbin4"});
    addRegions({"TR6Jbin1","TR6Jbin2","TR6Jbin3","TR6Jbin4"});

}

void OneLepton2018::ProcessEvent(AnalysisEvent *event)
{
    auto baseEle = event->getElectrons(7,2.47, ELooseBLLH);
	auto baseMuon = event->getMuons(6,2.70, MuMedium | MuNotCosmic | MuZ05mm | MuQoPSignificance);
	auto baseJets = event->getJets(20.,4.5);
	auto met_Vect  = event->getMET();
	float met       = met_Vect.Pt();
	auto weights        = event->getMCWeights();

    //
	// Baseline objects
	//
	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
	baseEle = overlapRemoval(baseEle, baseMuon, 0.01);
	// ele-jet OR
	baseJets = overlapRemoval(baseJets, baseEle, 0.2);
	baseEle = overlapRemoval(baseEle, baseJets, radiusCalcLep);
	// mu-jet OR
	baseJets = overlapRemoval(baseJets, baseMuon, 0.2, LessThan3Tracks);
	baseMuon = overlapRemoval(baseMuon, baseJets, radiusCalcLep);

	// 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 signalEle = filterObjects(baseEle,7., 2.47, ETightLH | ED0Sigma5  | EZ05mm  | EIsoFCLoose);
	AnalysisObjects signalEleLowPt;
    for (const auto& ele : signalEle) {
      if ((ele.Pt() < 200.)) signalEleLowPt.push_back(ele);
  	}
	auto signalEleHighPt = filterObjects(baseEle,200., 2.47, ETightLH | ED0Sigma5  | EZ05mm  | EIsoFCHighPtCaloOnly);
	signalEle = signalEleLowPt + signalEleHighPt;

	auto signalMuon = filterObjects(baseMuon,6., 2.7, MuD0Sigma3 | MuZ05mm | MuIsoFCLoose);
	auto signalLept =  signalEle + signalMuon;

	auto signalJets = filterObjects(baseJets, 30., 2.80, JVT120Jet);
	auto signalBJets = filterObjects(signalJets, 30., 2.8, BTag77MV2c10);

    unsigned int N_baseEle = baseEle.size();
    unsigned int N_baseMuon = baseMuon.size();
    unsigned int N_baseLept = N_baseEle + N_baseMuon;

    unsigned int N_signalEle = signalEle.size();
    unsigned int N_signalMuon = signalMuon.size();
    unsigned int N_signalLept = N_signalEle + N_signalMuon;
    unsigned int N_signalJets = signalJets.size();
    unsigned int N_signalBJets = signalBJets.size();

    // Preselection
    if(N_baseLept != 1) return;
    if(N_signalLept != 1) return;
    if(N_signalJets < 2) return;
    if(met< 240.) return;

    // Analysis observables
    float mt=0, ht=0, meff=0, apl=0, lepApl=0;
    mt   = calcMT(signalLept[0], met_Vect);
    ht   = sumObjectsPt(signalJets,999,30);
    meff = ht + met+ signalLept[0].Pt();
    apl   = aplanarity(signalJets);
    lepApl = aplanarity(signalJets + signalLept);

    ntupVar("AnalysisType", (baseMuon.size() == 1) ? 2 : 1);
    ntupVar("met",met);
    ntupVar("meff",meff);
    ntupVar("ht",ht);
    ntupVar("apl",apl);
    ntupVar("lepApl",lepApl);
    ntupVar("mt",mt);
    ntupVar("nJet30",(int)signalJets.size());
    ntupVar("nBJet30_MV2c10",(int)signalBJets.size());

    ntupVar("lep1Pt", signalLept[0].Pt());
    ntupVar("lep1Eta", signalLept[0].Eta());
    ntupVar("lep1Phi", signalLept[0].Phi());

    ntupVar("jet1Pt", signalJets[0].Pt());
    ntupVar("jet1Eta", signalJets[0].Eta());
    ntupVar("jet1Phi", signalJets[0].Phi());

    ntupVar("nLep_signal", (int)signalLept.size());
    ntupVar("nLep_base", (int)N_baseLept);


    if (signalLept[0].Pt()<25. && N_signalJets>=2 && met>400. && (met/meff)>0.25 && (N_signalJets/signalLept[0].Pt())>0.1) {
        // 2J regions
        if (mt>100.) {
            // 2J SRs
            if (N_signalBJets==0 && meff>700. && meff<=1300.) accept("SR2JBVbin1");
            else if (N_signalBJets==0 && meff>1300. && meff<=1900.) accept("SR2JBVbin2");
            else if (N_signalBJets==0 && meff>1900.) accept("SR2JBVbin3");
            else if (N_signalBJets>0 && meff>700. && meff<=1300.) accept("SR2JBTbin1");
            else if (N_signalBJets>0 && meff>1300. && meff<=1900.) accept("SR2JBTbin2");
            else if (N_signalBJets>0 && meff>1900.) accept("SR2JBTbin3");

            if (N_signalBJets==0 && meff>1900.) accept("SR2Jhigh");
            if (N_signalBJets==0 && meff>1300.) accept("SR2Jlow");
        }
        else if (mt>50. && mt<80.){
            // 2J CRs
            if (N_signalBJets==0 && meff>700. && meff<=1300.) accept("WR2Jbin1");
            else if (N_signalBJets==0 && meff>1300. && meff<=1900.) accept("WR2Jbin2");
            else if (N_signalBJets==0 && meff>1900.) accept("WR2Jbin3");
            else if (N_signalBJets>0 && meff>700. && meff<=1300.) accept("TR2Jbin1");
            else if (N_signalBJets>0 && meff>1300. && meff<=1900.) accept("TR2Jbin2");
            else if (N_signalBJets>0 && meff>1900.) accept("TR2Jbin3");
        }
    }
    else if (signalLept[0].Pt()>25. && N_signalJets>=4 && N_signalJets<6 && met>300. && lepApl>0.01 && (met/meff)>0.20) {
        // 4J regions
        if (mt>520.) {
            // 4J high-x SRs
            if (N_signalBJets==0 && meff>1000. && meff<=1600.) accept("SR4JhxBVbin1");
            else if (N_signalBJets==0 && meff>1600. && meff<=2200.) accept("SR4JhxBVbin2");
            else if (N_signalBJets==0 && meff>2200.) accept("SR4JhxBVbin3");
            else if (N_signalBJets>0 && meff>1000. && meff<=1600.) accept("SR4JhxBTbin1");
            else if (N_signalBJets>0 && meff>1600. && meff<=2200.) accept("SR4JhxBTbin2");
            else if (N_signalBJets>0 && meff>2200.) accept("SR4JhxBTbin3");

            if (N_signalBJets==0 && meff>2200.) accept("SR4Jhx");
        }
        else if (mt>150 && mt<520) {
            // 4J low-x SRs
            if (N_signalBJets==0 && meff>1000. && meff<=1600.) accept("SR4JlxBVbin1");
            else if (N_signalBJets==0 && meff>1600. && meff<=2200.) accept("SR4JlxBVbin2");
            else if (N_signalBJets==0 && meff>2200.) accept("SR4JlxBVbin3");
            else if (N_signalBJets>0 && meff>1000. && meff<=1600.) accept("SR4JlxBTbin1");
            else if (N_signalBJets>0 && meff>1600. && meff<=2200.) accept("SR4JlxBTbin2");
            else if (N_signalBJets>0 && meff>2200.) accept("SR4JlxBTbin3");

            if (N_signalBJets==0 && meff>2200.) accept("SR4Jlx");
        }
        else if (mt>50 && mt<90) {
            // 4J CRs
            if (N_signalBJets==0 && meff>1000. && meff<=1600.) accept("WR4Jbin1");
            else if (N_signalBJets==0 && meff>1600. && meff<=2200.) accept("WR4Jbin2");
            else if (N_signalBJets==0 && meff>2200.) accept("WR4Jbin3");
            else if (N_signalBJets>0 && meff>1000. && meff<=1600.) accept("TR4Jbin1");
            else if (N_signalBJets>0 && meff>1600. && meff<=2200.) accept("TR4Jbin2");
            else if (N_signalBJets>0 && meff>2200.) accept("TR4Jbin3");
        }
    }
    else if (signalLept[0].Pt()>25 && N_signalJets>=6 && lepApl>0.05) {
        // 6J regions
        if (met>300. && mt>225.) {
            //6J SRs
            if (N_signalBJets==0 && meff>700. && meff<=1400.) accept("SR6JBVbin1");
            else if (N_signalBJets==0 && meff>1400. && meff<=2100.) accept("SR6JBVbin2");
            else if (N_signalBJets==0 && meff>2100. && meff<=2800.) accept("SR6JBVbin3");
            else if (N_signalBJets==0 && meff>2800.) accept("SR6JBVbin4");
            else if (N_signalBJets>0 && meff>700. && meff<=1400.) accept("SR6JBTbin1");
            else if (N_signalBJets>0 && meff>1400. && meff<=2100.) accept("SR6JBTbin2");
            else if (N_signalBJets>0 && meff>2100. && meff<=2800.) accept("SR6JBTbin3");
            else if (N_signalBJets>0 && meff>2800.) accept("SR6JBTbin4");

            if (N_signalBJets==0 && meff>2800.) accept("SR6Jhigh");
            if (N_signalBJets==0 && meff>2100.) accept("SR6Jlow");
        }
        else if (met>250. && mt>50. && mt<100.) {
            // 6J CRs
            if (N_signalBJets==0 && meff>700. && meff<=1400.) accept("WR6Jbin1");
            else if (N_signalBJets==0 && meff>1400. && meff<=2100.) accept("WR6Jbin2");
            else if (N_signalBJets==0 && meff>2100. && meff<=2800.) accept("WR6Jbin3");
            else if (N_signalBJets==0 && meff>2800.) accept("WR6Jbin4");
            else if (N_signalBJets>0 && meff>700. && meff<=1400.) accept("TR6Jbin1");
            else if (N_signalBJets>0 && meff>1400. && meff<=2100.) accept("TR6Jbin2");
            else if (N_signalBJets>0 && meff>2100. && meff<=2800.) accept("TR6Jbin3");
            else if (N_signalBJets>0 && meff>2800.) accept("TR6Jbin4");
        }
    }

	return;
}
