#include "SimpleAnalysis/AnalysisClass.h"

DefineAnalysis(EwkThreeLeptonERJR2018)
void EwkThreeLeptonERJR2018::Init()
{
  addRegions({"Preselection"});
  addRegions({"SRlow","SRISR"});
  addRegions({"CRlow","CRISR"});
  addRegions({"VRlow","VRISR","VRISRsmallPTsoft","VRISRsmallRjetsinv"});
  //
  addHistogram("cutflow",20,0,20);
}

void EwkThreeLeptonERJR2018::ProcessEvent(AnalysisEvent *event)
{
   //====================================================================================================
   // get baseline objects
   // cfr. STconfig: https://gitlab.cern.ch/atlas-phys-susy-wg/AnalysisSUSYToolsConfigurations/blob/master/ANA-SUSY-2018-05.conf
   auto baselineElectrons  = event->getElectrons(10., 2.47, ELooseBLLH);
   auto baselineMuons      = event->getMuons(10, 2.4, MuMedium | MuNotCosmic | MuQoPSignificance);
   auto baselineJets       = event->getJets(20., 4.5);
   auto metVec             = event->getMET();
   double met              = metVec.Et();

   //====================================================================================================
   // overlap removal
   auto radiusCalcLepton = [] (const AnalysisObject& lepton, const AnalysisObject&) { return std::min(0.4, 0.04 + 10./lepton.Pt()); };
   // 1. e-mu: shared ID track
   baselineElectrons = overlapRemoval(baselineElectrons, baselineMuons, 0.01);
   // 2a. j-e: deltaR < 0.2
   baselineJets      = overlapRemoval(baselineJets, baselineElectrons, 0.2);
   // 2b. j-mu: deltaR < 0.2 but only LessThanThreeTrack
   baselineJets      = overlapRemoval(baselineJets, baselineMuons, 0.2, LessThan3Tracks);
   // 3. e/mu-j: (62.5:dR=0.2, 27.8: dR=0.4) variable radius: no leptons removed below/above 28/60
   baselineElectrons = overlapRemoval(baselineElectrons, baselineJets, radiusCalcLepton);
   baselineMuons     = overlapRemoval(baselineMuons, baselineJets, radiusCalcLepton);

   //====================================================================================================
   auto baselineleptons = baselineElectrons + baselineMuons;
   int nbaselineleptons = baselineleptons.size();

   //====================================================================================================
   // signal objects
   auto electrons = filterObjects(baselineElectrons, 20, 2.47, EMediumLH | ED0Sigma5 | EZ05mm | EIsoFCTight);
   auto muons     = filterObjects(baselineMuons, 20, 2.4, MuD0Sigma3 | MuZ05mm | MuIsoFCTightFR);
   auto jets      = filterObjects(baselineJets, 20, 2.4, JVT120Jet);
   auto bjets     = filterObjects(jets, 20., 2.4, BTag77MV2c10);
   auto leptons   = electrons + muons;

   //====================================================================================================
   int nleptons   = leptons.size();
   int nelectrons = electrons.size();
   int nmuons     = muons.size();
   int njets      = jets.size();
   int nbjets     = bjets.size();

   // pT order the objects
   sortObjectsByPt(leptons);
   sortObjectsByPt(jets);

   //====================================================================================================
   int cutflowcounter = 0;
   fill("cutflow", cutflowcounter++);

   // exactly 3 baseline and signal leptons
   if (nbaselineleptons != 3) return;
   if (nleptons != 3) return;
   fill("cutflow", cutflowcounter++);

   // b-jet veto
   if (nbjets>0) return;
   fill("cutflow", cutflowcounter++);

   //lepton pT cuts
   if ( (leptons[0].Pt()<25) && (leptons[1].Pt()<25) && (leptons[2].Pt()<20) ) return;
   fill("cutflow", cutflowcounter++);

   //====================================================================================================
   //do assignment by closest to mZ
   double mZ  = 91.1876;
   double mll = -999.;
   double mdiff = 1e6;
   int nSFOS = 0;
   int iZ1(-1), iZ2(-1), iW(-1);

   for (int ilep=0; ilep<nleptons-1; ilep++) {
      for (int jlep=ilep+1; jlep<nleptons; jlep++) {
         int klep = nleptons - ilep - jlep;
         if ( (leptons[ilep].type() == leptons[jlep].type()) && (leptons[ilep].charge() != leptons[jlep].charge()) ) {
            nSFOS++;
         //
            double imll = (leptons[ilep]+leptons[jlep]).M();
            double imdiff = fabs(mZ - imll);
         //
            if (imdiff < mdiff) {
               mdiff = imdiff;
               mll   = imll;
               iZ1   = ilep;
               iZ2   = jlep;
               iW    = klep;
            }
         }
      }
   }

   //====================================================================================================
   // remove events without SFOS
   if (nSFOS==0) return;
   fill("cutflow", cutflowcounter++);

   //====================================================================================================
   // at this point everything should be assigned!
   if (iZ1==-1 || iZ2==-1 || iW==-1) {
     std::cout << "Z(ll)W(lv) lepton assignment failed!" << std::endl;
   }

   //====================================================================================================
   // mlll > 105 GeV
   double mlll = (leptons[0]+leptons[1]+leptons[2]).M();
   if (mlll < 105.) return;
   fill("cutflow", cutflowcounter++);

   //====================================================================================================
   // require the Zpeak [75,105]
   if ( mll < 75. || mll > 105. ) return;
   fill("cutflow", cutflowcounter++);

   //====================================================================================================
   // require mTW
   double mTW = calcMT(leptons[iW], metVec);

   //====================================================================================================
   accept("Preselection");

   //====================================================================================================
   // additional variables:
   // ISR system
   TLorentzVector vISR = TLorentzVector(0,0,0,0);
   for (int ijet=0; ijet<(int)jets.size(); ++ijet) {
      TLorentzVector jet(0,0,0,0);
      jet.SetPtEtaPhiM(jets[ijet].Pt(), jets[ijet].Eta(), jets[ijet].Phi(), jets[ijet].M());
      vISR += jet;
   }
   // LEP system (massive leptons)
   TLorentzVector vLEP = TLorentzVector(0,0,0,0);
   for (int ilep=0; ilep<(int)leptons.size(); ++ilep) {
      TLorentzVector lep(0,0,0,0);
      lep.SetPtEtaPhiM(leptons[ilep].Pt(), leptons[ilep].Eta(), leptons[ilep].Phi(), leptons[ilep].M());
      vLEP += lep;
   }

   TLorentzVector vMET = TLorentzVector(0,0,0,0);
   vMET.SetPxPyPzE(metVec.Px(), metVec.Py(), 0., sqrt(metVec.Px()*metVec.Px()+metVec.Py()*metVec.Py()));

   double pTjets(-999), dphijetsinv(-999), Rjetsinv(-999), pTsoft(-999);
   if (jets.size()>0) {
      pTjets      = vISR.Pt();
      dphijetsinv = vISR.DeltaPhi(vMET);
      Rjetsinv    = fabs(vMET.Px()*vISR.Px() + vMET.Py()*vISR.Py()) / fabs(pTjets*pTjets);
      pTsoft      = (vLEP + vMET + vISR).Pt();
   }

   double pTsoftPP(-999), meff3l(-999), long_inv(-999), pI(-999), boostx(-999), boosty(-999), boostz(-999);
   pTsoftPP = (vLEP + vMET).Pt();
   meff3l   = leptons[0].Pt()+ leptons[1].Pt() + leptons[2].Pt() + met;

   //calculating Z-component of the missing energy vector
   long_inv = vLEP.Pz()*sqrt(met*met) / sqrt(vLEP.Pt()*vLEP.Pt()+vLEP.M()*vLEP.M());
   pI = sqrt(vMET.Px()*vMET.Px() + vMET.Py()*vMET.Py() + long_inv*long_inv);

   //calculating the boost
   boostx = (vLEP.Px() + vMET.Px()) / (vLEP.E() + sqrt(pI*pI));
   boosty = (vLEP.Py() + vMET.Py()) / (vLEP.E() + sqrt(pI*pI));
   boostz = (vLEP.Pz() + long_inv) / (vLEP.E() + sqrt(pI*pI));

   TLorentzVector vMETprime;
   vMETprime.SetXYZM(vMET.Px(), vMET.Py(), long_inv, 0.);
   //
   std::vector<TLorentzVector> leptons_boost(3,TLorentzVector(0.,0.,0.,0.));
   for (int ilep=0; ilep<3; ilep++) {
      leptons_boost[ilep].SetPtEtaPhiM(leptons[ilep].Pt(), leptons[ilep].Eta(), leptons[ilep].Phi(), leptons[ilep].M());
      leptons_boost[ilep].Boost(-boostx, -boosty, -boostz);
   }

   double H_boost(-999), HTratio(-999), pTratio(-999);
   H_boost  = leptons_boost[0].P() + leptons_boost[1].P() + leptons_boost[2].P() + vMETprime.P();

   HTratio = meff3l / H_boost;
   pTratio = pTsoftPP / (pTsoftPP + meff3l);

   //====================================================================================================
   // CR/VR/SR regions: low / ISR / ISRadditional
   if (leptons[0].Pt()>60. && leptons[1].Pt()>40. && leptons[2].Pt()>30. && njets==0 && mTW>0.   && mTW<70.  && H_boost>250. && pTratio<0.2  && HTratio>0.75 && met>40.) accept("CRlow");
   if (leptons[0].Pt()>60. && leptons[1].Pt()>40. && leptons[2].Pt()>30. && njets==0 && mTW>70.  && mTW<100. && H_boost>250. && pTratio<0.2  && HTratio>0.75)            accept("VRlow");
   if (leptons[0].Pt()>60. && leptons[1].Pt()>40. && leptons[2].Pt()>30. && njets==0 && mTW>100. &&             H_boost>250. && pTratio<0.05 && HTratio>0.9)             accept("SRlow");
   //
   if (leptons[0].Pt()>25. && leptons[1].Pt()>25. && leptons[2].Pt()>20. && njets>0            && mTW<100. && fabs(dphijetsinv)>2.0 && Rjetsinv>0.55 && Rjetsinv<1.0  && pTjets>80.  && met>60. && pTsoft<25.) accept("CRISR");
   if (leptons[0].Pt()>25. && leptons[1].Pt()>25. && leptons[2].Pt()>20. && njets>0            && mTW>60.  && fabs(dphijetsinv)>2.0 && Rjetsinv>0.55 && Rjetsinv<1.0  && pTjets>80.  && met>60. && pTsoft>25.) accept("VRISR");
   if (leptons[0].Pt()>25. && leptons[1].Pt()>25. && leptons[2].Pt()>20. && njets>0 && njets<4 && mTW>100. && fabs(dphijetsinv)>2.0 && Rjetsinv>0.55 && Rjetsinv<1.0  && pTjets>100. && met>80. && pTsoft<25.) accept("SRISR");
   //
   if (leptons[0].Pt()>25. && leptons[1].Pt()>25. && leptons[2].Pt()>20. && njets>0            && mTW>60.  && fabs(dphijetsinv)>2.0 && Rjetsinv>0.55 && Rjetsinv<1.0  && pTjets<80.  && met>60. && pTsoft<25.) accept("VRISRsmallPTsoft");
   if (leptons[0].Pt()>25. && leptons[1].Pt()>25. && leptons[2].Pt()>20. && njets>0            && mTW>60.  && fabs(dphijetsinv)>2.0 && Rjetsinv>0.3  && Rjetsinv<0.55 && pTjets>80.  && met>60. && pTsoft<25.) accept("VRISRsmallRjetsinv");
   //

   ntupVar("met", met);
   ntupVar("mll", mll);
   ntupVar("mTW", mTW);
   ntupVar("mlll", mlll);
   //
   ntupVar("njets", njets);
   ntupVar("nbjets", nbjets);
   ntupVar("nbaselineleptons", nbaselineleptons);
   ntupVar("nleptons", nleptons);
   ntupVar("nelectrons", nelectrons);
   ntupVar("nmuons", nmuons);
   ntupVar("nSFOS", nSFOS);
   //
   ntupVar("lep1", leptons[0]);
   ntupVar("lep2", leptons[1]);
   ntupVar("lep3", leptons[2]);
   ntupVar("lepZ1", leptons[iZ1]);
   ntupVar("lepZ2", leptons[iZ2]);
   ntupVar("lepW", leptons[iW]);
   // ISR
   ntupVar("H_boost", H_boost);
   ntupVar("HTratio", HTratio);
   ntupVar("pTratio", pTratio);
   ntupVar("dphijetsinv", dphijetsinv);
   ntupVar("Rjetsinv", Rjetsinv);
   ntupVar("pTjets", pTjets);
   ntupVar("pTsoft", pTsoft);
   // low
   ntupVar("pTsoftPP", pTsoftPP);
   ntupVar("meff3l", meff3l);
   //
   ntupVar("susyProcess", event->getSUSYChannel());
   ntupVar("mcDSID", event->getMCNumber());
   ntupVar("mcWeights", event->getMCWeights());

   return;
}
