#include "SimpleAnalysis/AnalysisClass.h"

DefineAnalysis(EwkThreeLeptonResonance2018)

void EwkThreeLeptonResonance2018::Init()
{

  addRegions({"CRWZ", "CRZj", "CRttZ", "CRZZ"});
  addRegions({"VRMet", "VRmTmin", "VRZj", "VRttZ", "VRZZ"});
  addRegions({"SR3l", "SR4l", "SRFR"});

  std::vector<int> mZl_lowedges {90,110,130,150,170,190,210,230,250,270,300,330,360,400,440,580};
  for(unsigned int i=0; i < mZl_lowedges.size(); i++){
    addRegions( {"SR3l_"+std::to_string(mZl_lowedges[i])} );
  }
  for(unsigned int i=0; i < mZl_lowedges.size(); i++){
    addRegions( {"SR4l_"+std::to_string(mZl_lowedges[i])} );
  }
  for(unsigned int i=0; i < mZl_lowedges.size(); i++){
    addRegions( {"SRFR_"+std::to_string(mZl_lowedges[i])} );
  }
  
}

std::pair<int,int> IdentifyZleptons(AnalysisObjects objects, bool doSFOS=false)
{
  //Find two objects (leptons/jets) closest to Z mass
  auto object_index = std::make_pair(-1, -1);
  float Zmass = 999999;
  float this_Zmass = 0;
  for( unsigned int i=0; i < objects.size(); i++){
    for( unsigned int j=i+1; j < objects.size(); j++){

      if(doSFOS && ((objects[i].charge() == objects[j].charge()) || (objects[i].type() != objects[j].type())) )
          continue;

      this_Zmass = (objects[i] + objects[j]).M(); 

      if( std::abs(this_Zmass-91.2) < std::abs(Zmass-91.2) ){
        Zmass = this_Zmass;
        object_index = std::make_pair(i,j);
      }

    }
  }
  return object_index;
}

bool doesSFOSexist(AnalysisObjects leptons)
{
  bool found_SFOS_pair = false;
  for( unsigned int i=0; i < leptons.size(); i++){
    for( unsigned int j=i+1; j < leptons.size(); j++){
      if( (leptons[i].charge() != leptons[j].charge()) && (leptons[i].type() == leptons[j].type()) ){
        found_SFOS_pair = true;
        break;
      }
    }//for j
  }//for i
  return found_SFOS_pair;
}

bool doesSFexist(AnalysisObjects leptons)
{
  bool found_SF_pair = false;
  for( unsigned int i=0; i < leptons.size(); i++){
    for( unsigned int j=i+1; j < leptons.size(); j++){
      if( leptons[i].type() == leptons[j].type() ){
        found_SF_pair = true;
        break;
      }
    }//for j
  }//for i
  return found_SF_pair;
}

void EwkThreeLeptonResonance2018::ProcessEvent(AnalysisEvent *event)
{
  
  std::vector<float> mZl_lowedges {90,110,130,150,170,190,210,230,250,270,300,330,360,400,440,580};

  //SUSY2 derivation & trigger selection
  auto susy2_el = event->getElectrons(9, 2.6, ELooseLH);
  auto susy2_mu = event->getMuons(9, 2.6, MuLoose);
  sortObjectsByPt(susy2_el);
  sortObjectsByPt(susy2_mu);

  if (susy2_el.size()+susy2_mu.size() < 2 )
    return;

  //Trigger pt threshold requirement 
  if ( !((susy2_el.size() > 0 && susy2_el.at(0).Pt() >= 26) || (susy2_mu.size() > 0 && susy2_mu.at(0).Pt() > 26)) )
    return;


  //Baseline objects
  auto baseline_jet = event->getJets(20., 4.8);  /// JVT cut not applied
  auto baseline_el  = filterCrack(event->getElectrons(10, 2.47, ELooseBLLH|EZ05mm));
  auto baseline_mu  = event->getMuons(10, 2.7, MuMedium|MuZ05mm);
  auto MET_vec      = event->getMET();
  float MET         = MET_vec.Et();


  //Half-HF overlap removal
  auto radiusCalcBjetLep   = [] (const AnalysisObject& bjet, const AnalysisObject& ) { return 0.2*(bjet.Pt() > 100); };

  baseline_el  = overlapRemoval(baseline_el,  baseline_mu,  0.01);
  baseline_jet = overlapRemoval(baseline_jet, baseline_el,  0.2, NOT(BTag85MV2c10));
  baseline_jet = overlapRemoval(baseline_jet, baseline_el,  radiusCalcBjetLep, BTag85MV2c10);
  baseline_el  = overlapRemoval(baseline_el,  baseline_jet, 0.4);
  baseline_jet = overlapRemoval(baseline_jet, baseline_mu,  0.2, LessThan3Tracks|NOT(BTag85MV2c10));
  baseline_jet = overlapRemoval(baseline_jet, baseline_mu,  radiusCalcBjetLep, LessThan3Tracks|BTag85MV2c10);
  baseline_mu  = overlapRemoval(baseline_mu,  baseline_jet, 0.4);
  
  // Remove events with one bad jet
  if (countObjects(baseline_jet, 20, 2.8, NOT(LooseBadJet)) != 0) return;

  // Signal objects
  auto signal_jet  = filterObjects(baseline_jet, 20, 2.8, JVT120Jet|LooseBadJet); 
  auto signal_bjet = filterObjects(signal_jet,   20, 2.5, BTag85MV2c10);
  auto signal_el   = filterObjects(baseline_el,  12, 2.47, EMediumLH|EIsoFCTight|EZ05mm|ED0Sigma5);  
  auto signal_mu   = filterObjects(baseline_mu,  12, 2.7, MuD0Sigma3|MuZ05mm|MuMedium|MuIsoFCTightFR); 

  sortObjectsByPt(signal_jet);
  sortObjectsByPt(signal_bjet);
  sortObjectsByPt(signal_el);
  sortObjectsByPt(signal_mu);
  
  auto signal_lep = signal_el+signal_mu;
  AnalysisClass::sortObjectsByPt(signal_lep);
  auto baseline_lep = baseline_el+baseline_mu;
  AnalysisClass::sortObjectsByPt(baseline_lep);

  //Preselection
  if(signal_lep.size() < 3){
    return;
  }
    
  //Calculate LT of the event  
  float LT = 0;
  for(unsigned int i=0; i < signal_lep.size(); ++i){
    LT += signal_lep[i].Pt();
  }

  //Find two leptons closest to Z mass
  auto Z1_lep_index = IdentifyZleptons(signal_lep, /*doSFOS*/ true);
  if(Z1_lep_index.first == -1){
    return; //No SFOS pair found
  }
  AnalysisObject Z1_4vec = signal_lep[Z1_lep_index.first]+signal_lep[Z1_lep_index.second];
  float Z1_mass = (signal_lep[Z1_lep_index.first] + signal_lep[Z1_lep_index.second]).M(); 

  if( std::abs(Z1_mass - 91.2) > 10 )
    return;

  auto remaining_lep = AnalysisObjects(signal_lep);
  remaining_lep.erase( remaining_lep.begin()+Z1_lep_index.second );
  remaining_lep.erase( remaining_lep.begin()+Z1_lep_index.first );
  
  //Check for a second Z in 4-lepton events (ZZ background)
  float other_Z_mass_4lep = -1;
  if( signal_lep.size() == 4 ){
    if( doesSFOSexist(remaining_lep) )
      other_Z_mass_4lep = (remaining_lep[0] + remaining_lep[1]).M();
  }
  
  //Does event have an additional same-flavor pair (on top of first Z pair), for SR4l MET requirement
  bool has2SFpairs = false;
  if( doesSFexist(remaining_lep) )
    has2SFpairs = true;

  //Find a second signal boson (not necessarily a Z)
  std::pair<int,int> Z2_obj_index; 
  float Z2_mass = -1;
  bool is_lep = false;

  //First check for second signal Z from remaining leptons
  if( remaining_lep.size() >=4 ){
    Z2_obj_index = IdentifyZleptons(remaining_lep, /*doSFOS*/ true); 
    if( Z2_obj_index.first != -1 ){
      //Find position in original lepton vector
      int index_lep1 = std::find(signal_lep.begin(), signal_lep.end(), remaining_lep[Z2_obj_index.first]  ) - signal_lep.begin();
      int index_lep2 = std::find(signal_lep.begin(), signal_lep.end(), remaining_lep[Z2_obj_index.second] ) - signal_lep.begin();
      Z2_obj_index = std::make_pair(index_lep1, index_lep2);
      Z2_mass = (signal_lep[Z2_obj_index.first] + signal_lep[Z2_obj_index.second]).M();
      is_lep = true; //2nd Z is leptonic
    }

    //If not close to Zmass
    if( std::abs(Z2_mass-91.2) > 10 ){
      Z2_mass = -1;
      Z2_obj_index = std::make_pair(-1, -1 );
      is_lep = false; //2nd Z is not leptonic
    }
  } 

  //Then check for second signal Z / W / H from jets
  if( remaining_lep.size() >= 2 && signal_jet.size() >= 2 ){
    std::pair<int,int> this_Z2_obj_index = IdentifyZleptons(signal_jet, /*doSFOS*/ false); 
    if( this_Z2_obj_index.first != -1 ){
      float this_Z2_mass = (signal_jet[this_Z2_obj_index.first] + signal_jet[this_Z2_obj_index.second]).M();
    
      bool hasBjet = false;
      int this_mass_diff = 999999;
      if( signal_jet[this_Z2_obj_index.first].pass(BTag85MV2c10) || signal_jet[this_Z2_obj_index.second].pass(BTag85MV2c10) ){
        hasBjet = true;
        this_mass_diff = std::abs(this_Z2_mass - 125.);
      } else {
        this_mass_diff = std::abs(this_Z2_mass - 91.2);
      }
      
      //If close enough to Zmass
      if( (hasBjet && this_Z2_mass > 71.2 && this_Z2_mass < 150) || (!hasBjet && this_Z2_mass > 71.2 && this_Z2_mass < 111.2) ){
        //If we found Z2 mass from remaining leptons, but this jet one is closer to Zmass, choose this one
        if(Z2_mass == -1 || (this_mass_diff < std::abs(Z2_mass-91.2)) ){ 
          Z2_obj_index = this_Z2_obj_index;
          Z2_mass = this_Z2_mass;
          is_lep = false; //2nd Z is not leptonic
        }
      }
    }
  }

  AnalysisObject C1_vec(-1, 0, 0, 0, 0, 0, NONE, 0, 0 );
  AnalysisObject C2_vec(-1, 0, 0, 0, 0, 0, NONE, 0, 0 );

  //If we have the 2nd boson, find lepton pairings to give EWKinos
  float mZl_asym = 999999;
  std::pair<int,int> EWKino_lep_index = std::make_pair(-1, -1); 
  if( Z2_mass != -1 ){

    AnalysisObject Z2_4vec( -1, 0, 0, 0, 0, 0, NONE, 0, 0);
    if(is_lep){
      Z2_4vec = signal_lep[Z2_obj_index.first]+signal_lep[Z2_obj_index.second];
      //Remove the used leptons if the 2nd boson was a lepton Z
      int index_lep1 = std::find(remaining_lep.begin(), remaining_lep.end(), signal_lep[Z2_obj_index.first]  ) - remaining_lep.begin();
      int index_lep2 = std::find(remaining_lep.begin(), remaining_lep.end(), signal_lep[Z2_obj_index.second] ) - remaining_lep.begin();
      if( index_lep1 > index_lep2 ){
        int tmp = index_lep2;
        index_lep2 = index_lep1;
        index_lep1 = tmp;
      }
      remaining_lep.erase( remaining_lep.begin()+index_lep2 );
      remaining_lep.erase( remaining_lep.begin()+index_lep1 );
    } else {
      Z2_4vec = signal_jet[Z2_obj_index.first]+signal_jet[Z2_obj_index.second];
    }

    float this_mZl_asym = -1;
    for(unsigned int i=0; i < remaining_lep.size(); ++i){
      for(unsigned int j=0; j < remaining_lep.size(); ++j){
        if( i == j )
          continue;
        float mass_pair1 = (Z1_4vec+remaining_lep[i]).M();
        float mass_pair2 = (Z2_4vec+remaining_lep[j]).M();
        this_mZl_asym = std::abs(mass_pair1-mass_pair2)/(mass_pair1+mass_pair2);
        if( this_mZl_asym < mZl_asym ){
          mZl_asym = this_mZl_asym;
          C1_vec = Z1_4vec+remaining_lep[i];
          C2_vec = Z2_4vec+remaining_lep[j];

          int index_lep1 = std::find(signal_lep.begin(), signal_lep.end(), remaining_lep[i]  ) - signal_lep.begin();
          int index_lep2 = std::find(signal_lep.begin(), signal_lep.end(), remaining_lep[j] ) - signal_lep.begin();
          EWKino_lep_index = std::make_pair(index_lep1,index_lep2);
        }
      }//for j
    }//for i
  }
  else { //If we don't have a Z2


    if( LT < 550 ){ //Choose closest dR
      float min_dR = 999.;
      float this_min_dR = 999.;
      for(unsigned int i=0; i < remaining_lep.size(); ++i){
        this_min_dR = Z1_4vec.DeltaR( remaining_lep[i] );
        if( this_min_dR < min_dR ){
          min_dR = this_min_dR;
          C1_vec = Z1_4vec + remaining_lep[i];
          int index_lep1 = std::find(signal_lep.begin(), signal_lep.end(), remaining_lep[i]  ) - signal_lep.begin();
          EWKino_lep_index = std::make_pair( index_lep1, -1 );
        }
      }
    } else { //LT > 550
      float max_mass = -1.;
      float this_max_mass = -1.;
      for(unsigned int i=0; i < remaining_lep.size(); ++i){
        this_max_mass = (Z1_4vec + remaining_lep[i]).M();
        if( this_max_mass > max_mass ){
          max_mass = this_max_mass;
          C1_vec = Z1_4vec + remaining_lep[i];
          int index_lep1 = std::find(signal_lep.begin(), signal_lep.end(), remaining_lep[i]  ) - signal_lep.begin();
          EWKino_lep_index = std::make_pair( index_lep1, -1 );
        }
      }
    }//LT > 550
  }

    
  //Calculate mTmin of the event.  Lowest mTmin, requiring remaining leptons to have a SFOS pair
  float mTmin = 9999.;
  float this_mTmin = 0;
  for(unsigned int i=0; i < signal_lep.size(); ++i){
    auto other_lep = AnalysisObjects(signal_lep); //vector of signal leptons without this one
    other_lep.erase( other_lep.begin()+i );

    if( doesSFOSexist(other_lep) ){
      this_mTmin = sqrt(2*signal_lep[i].Pt()*MET_vec.Pt()*(1-TMath::Cos( signal_lep[i].DeltaPhi(MET_vec) )));
      mTmin = std::min(mTmin, this_mTmin);
    }
  }

  //Calcluate dRbb of an event if there are 2 bjets
  float dRbb = 0;
  if( signal_bjet.size() >= 2 ){
    dRbb = signal_bjet[0].DeltaR(signal_bjet[1]);    
  }

  if (signal_lep.size() == 3 && dRbb < 1.5){
    
    if( MET > 150 && mTmin > 125 ){
      if( C1_vec.M() < 90 )
        return;
      accept("SR3l");
      int i_mZl_bin = std::lower_bound(mZl_lowedges.begin(), mZl_lowedges.end(), C1_vec.M() ) - mZl_lowedges.begin() - 1;
      accept( "SR3l_"+std::to_string(int(mZl_lowedges[i_mZl_bin])) );
    } else if( MET < 80 && mTmin > 50 && mTmin < 100 ){
      accept("CRWZ");
    } else if( MET > 80 && mTmin < 100 ){
      accept("VRMet");
    } else if( MET < 80 && mTmin > 125 ){
      accept("VRmTmin");
    } else if( MET < 30 && mTmin < 30 ){
      accept("CRZj");
    } else if( MET > 30 && MET < 80 && mTmin < 30 ){
      accept("VRZj");
    }


  } else if( dRbb > 1.5 ){

    //Do not allow a 2nd Z in 4 lepton events
    if( other_Z_mass_4lep != -1 && std::abs(other_Z_mass_4lep-91.2) < 20 )
      return;

    if( dRbb > 2.5 && MET > 40 ){
      accept("CRttZ");
    } else if( dRbb < 2.5 && MET > 40 ){
      accept("VRttZ");
    }

  } else if( signal_lep.size() >= 4 ){

    if( std::abs(other_Z_mass_4lep-91.2) < 5 ){
      accept("CRZZ");
      return;
    } else if (std::abs(other_Z_mass_4lep-91.2) < 20 ){
      accept("VRZZ");
      return;
    }
    
    if ( Z2_mass != -1 && mZl_asym < 0.1) {
      float this_EWKino_mass = (C1_vec.M()+C2_vec.M())/2.;
      if( this_EWKino_mass < 90. )
        return;

      accept("SRFR");
      int i_mZl_bin = std::lower_bound(mZl_lowedges.begin(), mZl_lowedges.end(), this_EWKino_mass ) - mZl_lowedges.begin() - 1;
      accept( "SRFR_"+std::to_string(int(mZl_lowedges[i_mZl_bin])) );

      return;
    }
    
    if ( Z2_mass == -1 && (!has2SFpairs || MET > 80) ) {
      if( C1_vec.M() < 90 )
        return;
      accept("SR4l");
      int i_mZl_bin = std::lower_bound(mZl_lowedges.begin(), mZl_lowedges.end(), C1_vec.M() ) - mZl_lowedges.begin() - 1;
      accept( "SR4l_"+std::to_string(int(mZl_lowedges[i_mZl_bin])) );
      return;
    }
  }
}
