//this is the main script that I ran when conducting trigger efficiency analysis. however, there are other scripts, far longer. that do physics analysis that are used to do a further trigger analysis.

#include <tea/ComputeEffAlg.h>
#include "TrigSteeringEvent/TrigRoiDescriptorCollection.h"
#include <xAODEventInfo/EventInfoContainer.h>
#include <xAODTruth/TruthParticleContainer.h>
#include <xAODTruth/TruthVertexContainer.h>

ComputeEffAlg :: ComputeEffAlg (const std::string& name,
                                  ISvcLocator *pSvcLocator): EL::AnaAlgorithm (name, pSvcLocator)
{
  declareProperty("trigger_names", m_trigger_names = "", "Space seperated trigger chain names");
  declareProperty("combined_trigger_names", m_combined_trgger_names = "", "Space seperated combined trigger chain names");
  declareProperty("jet_collection_name", m_jet_collection_key = "HLT_AntiKt4EMTopoJets_subjesIS", "Name of jet collection");
  declareProperty("enable_truth", m_enable_truth = false, "Enable truth");
  declareProperty("truth_pdgid", m_truth_pdgid=-1, "Truth PDG ID");
}

StatusCode ComputeEffAlg::initialize(){
  m_all = m_trigDec->getChainGroup( "HLT_.*" );

  parse_names();

  for(auto pair: probed_triggers){
    TString tn(pair.first);
//book the needed histograms 
    ANA_CHECK(book(TH1F("leading_jet_pt_"+tn, "p_{T} of the leading jet for trigger "+tn+"; p_{T} [GeV]", 80, 0, 800)));
    ANA_CHECK(book(TH1F("mjj_"+tn, "MJJ for Trigger "+tn+"; M_{JJ} [GeV]", 80, 0, 1500)));
    hist("mjj_"+pair.first)->Sumw2();
    hist("leading_jet_pt_"+pair.first)->Sumw2();

    ANA_CHECK(book(TH1F("kin_leading_jet_pt_"+tn, "p_{T} of the leading jet for trigger "+tn+"; p_{T} [GeV]", 80, 0, 800)));
    ANA_CHECK(book(TH1F("kin_mjj_"+tn, "MJJ for Trigger "+tn+"; M_{JJ} [GeV]", 80, 0, 1500)));
    hist("kin_mjj_"+pair.first)->Sumw2();
    hist("kin_leading_jet_pt_"+pair.first)->Sumw2();

    if(m_enable_truth){
      ANA_CHECK(book(TH1F("truth_decayradius_"+tn, "Truth decay vtx for trigger "+tn+"; min vertex decay radius [GeV]", 400, 0, 800)));
    }

    for(auto x : pair.second){
      ANA_MSG_INFO("Combined trigger "<<pair.first<<" -> "<<x);
    }
  }

  ANA_MSG_INFO("Jet container key = '"<<m_jet_collection_key<<"'");

  return StatusCode::SUCCESS;
}


StatusCode ComputeEffAlg::execute(){
  const xAOD::JetContainer* jets = nullptr;
  const xAOD::TruthParticleContainer* truth_particles = nullptr;
  const xAOD::EventInfo* evt_info = nullptr;

  ANA_CHECK (evtStore() -> retrieve (evt_info, "EventInfo"));
  ANA_CHECK (evtStore() -> retrieve (jets, m_jet_collection_key));

  if(m_enable_truth){
    ANA_CHECK (evtStore() -> retrieve (truth_particles, "TruthParticles"));
  }

  //compute quantities for each event
  double leading_pt = -999;
  double subleading_pt = -999;
  const xAOD::Jet* leading_jet = nullptr;
  const xAOD::Jet* subleading_jet = nullptr;
  for(auto j: *jets){
    if(j->pt()/1000.0 > leading_pt){
      leading_pt = j->pt()/1000.0;
      leading_jet = j;
    }
  }
  std::vector<const xAOD::Jet*> pt_sorted_jets;
  for(auto j: *jets){
    if(j->pt()/1000.0 > subleading_pt && j->pt()/1000.0 < leading_pt){
      subleading_pt = j->pt()/1000.0; 
      subleading_jet = j;
    }

    pt_sorted_jets.push_back(j);
  }


  if(leading_jet == nullptr || subleading_jet == nullptr){
    ANA_MSG_INFO("Missing Jet");
    return StatusCode::SUCCESS;
  }

  //sort jets by pt
  std::sort(pt_sorted_jets.begin(), pt_sorted_jets.end(), [](const xAOD::Jet* a, const xAOD::Jet* b) -> bool
    {
        return a->pt() > b->pt();
    }
  );

//loop over pt_sorted
  double mjj = -619;
  const xAOD::Jet* mjj_j1 = nullptr;
  const xAOD::Jet* mjj_j2 = nullptr;
  for(auto i: pt_sorted_jets){
    if(std::abs(i->eta()) > 3.2 || i->pt()/1000 < 70) continue;
    for(auto j: pt_sorted_jets){
      if(std::abs(j->eta()) > 4.9 || j->pt()/1000 < 50 || i == j) continue;
      if(std::abs(j->eta() - i->eta()) < 0.4 || std::abs((j->p4().DeltaPhi(i->p4()) > 2))) continue;
      double mij = std::abs((i->p4()+j->p4()).M())/1000;
      if(mjj < mij){
        mjj_j1 = i;
        mjj_j2 = j;
        mjj = mij;
        }
    }
  }

  //truth decay vtx
  double truth_decayvtx = -100;

  if(m_enable_truth){
    double min=80000;
    for(auto tp: *truth_particles){
      if(tp->absPdgId() == m_truth_pdgid){
        if(!tp->hasDecayVtx()) continue;
        //calculate the radius for this particle
        double r = std::abs(tp->decayVtx()->perp());

        if(r<min){
          min = r;
        }
      }
    }

    truth_decayvtx = min;
  }
double weight = evt_info->mcEventWeight();
bool pass_trig_kin = false;

int okay_trigger_jets = 0;

for(auto j:*jets){
  if (j->pt()/1000 < 50) continue;
  if(std::abs(j->eta()) > 2.4) continue;
  okay_trigger_jets += 1;
}
pass_trig_kin = (okay_trigger_jets >= 2);

  for(auto pair: probed_triggers){
    bool combined_pass(false);

    for(auto trig_name : pair.second){
      if(trig_name == "AcceptAll"){
        combined_pass = true;
      }
      if(m_trigDec->isPassed(trig_name)){
        combined_pass=true;
      }
    }

    if(combined_pass){
      hist("leading_jet_pt_"+pair.first)->Fill(leading_pt, weight);
      hist("mjj_"+pair.first)->Fill(mjj, weight);
      if(m_enable_truth){
        hist("truth_decayradius_"+pair.first)->Fill(truth_decayvtx);
      }
      if(pass_trig_kin){
        hist("kin_mjj_"+pair.first)->Fill(mjj, weight);
        hist("kin_leading_jet_pt_"+pair.first)->Fill(leading_pt, weight);
      }
    }
  }

  return StatusCode::SUCCESS;
}


StatusCode ComputeEffAlg::finalize(){
  return StatusCode::SUCCESS;
}

void ComputeEffAlg::parse_names(){
  size_t last_index = 0;

  while(last_index != std::string::npos){
    size_t new_index = m_trigger_names.find(" ", last_index);

    std::string ss = m_trigger_names.substr(last_index, new_index-last_index);

    if(ss.empty()) break;

    probed_triggers[ss] = std::vector<std::string>{ss};

    if(new_index == std::string::npos){break;}

    last_index = new_index+1;
  }

  last_index = 0;

  while(last_index != std::string::npos){
    size_t new_index = m_combined_trgger_names.find(" ", last_index);

    std::string ss = m_combined_trgger_names.substr(last_index, new_index-last_index);

    if(ss.empty()) break;

    //split into parts
    size_t parts_last_index = 0;
    std::string pname="";
    std::vector<std::string> parts;

    while(parts_last_index != std::string::npos){
      size_t ni = ss.find(":", parts_last_index);
      std::string ps = ss.substr(parts_last_index, ni-parts_last_index);

      if(pname.empty()){
        pname = ps;
        continue;
      }

      parts.push_back(ps);

      if(ni == std::string::npos) break;

      parts_last_index = ni+1;
    }

    probed_triggers[pname] = parts;

    if(new_index == std::string::npos){break;}

    last_index = new_index+1;
  }
}