// -*- C++ -*-
//#include "HepMC/GenParticle.h"
#include "HepMC3/GenParticle.h"
#include "Rivet/Analysis.hh"
#include "Rivet/Projections/FinalState.hh"
#include "Rivet/Projections/FastJets.hh"
#include <regex>

//using HepMC::GenParticle;
using HepMC3::GenParticle;


namespace Rivet {
  // a work based on VBF HWW analysis and Rivet routine from Higgs differential cross-section combination between the ATLAS measurements in the yy and 4l channels
  //https://gitlab.com/hepcedar/rivet/-/blob/master/analyses/pluginATLAS/MY_ANALYSIS.cc
  //
  /// @author Jiayi Chen <jiayi.chen@cern.ch>
  class MY_ANALYSIS : public Analysis {
  public:

    /// Constructor
    //RIVET_DEFAULT_ANALYSIS_CTOR(MY_ANALYSIS);
    DEFAULT_RIVET_ANALYSIS_CTOR(MY_ANALYSIS);


    /// Book histograms and initialise projections before the run
    void init() {

      // All final state particles
      const FinalState fs;
      declare(fs, "FS");
      
      // Histograms 
      vector <double> mjjBinning;
      for(double x=0; x<1.1e3; x+=200.) mjjBinning.push_back(x);
      mjjBinning.push_back(1.2e3);
      vector <double> Mjj__signDPhijjBinningX;
      for(double x=300; x<1.1e3; x+=(1.1e3-300)/196 ){  Mjj__signDPhijjBinningX.push_back(x); }
            
      for(std::vector<string>::iterator it=selectionSetup.begin(); it!=selectionSetup.end(); it++){
	book(_h["xsec"+(*it)]      , "xsec"+(*it),1,  0,   1);
	book(_h["Njets"+(*it)]      , "Njets"+(*it),5,  0,   5);
	book(_h["pT_H"+(*it)]   , "pT_H"+(*it),{0,80,120,260,500,850});
	book(_h["Mjj"+(*it)], "Mjj"+(*it),mjjBinning);
	book(_h["sign_DPhijj"+(*it)], "sign_DPhijj"+(*it),
	     {-M_PI,-0.5*M_PI,0,0.5*M_PI,M_PI});
	book(_h["DYjj"+(*it)], "DYjj"+(*it),{0,2,4,5,6,7,8});
	book(_h["pTHjj"+(*it)], "pTHjj"+(*it),100,0,2e3);
	book(_h2D["Mjj__signDPhijj"+(*it)], "Mjj__signDPhijj"+(*it),Mjj__signDPhijjBinningX,
	     {-M_PI,-0.5*M_PI,0,0.5*M_PI,M_PI});
	book(_h2D["Mjj__pTH"+(*it)], "Mjj__pTH"+(*it),mjjBinning,
	     {0,80,120,260,500,850,10000});
	book(_h2D["DYjj__pTH"+(*it)], "DYjj__pTH"+(*it),{0,2,9},
	     {0,80,120,260,500,850,10000});
	book(_h2D["Njet__pTH"+(*it)], "Njet__pTH"+(*it),{0,1,2,3,4,5},
	     {0,80,120,260,500,850,10000});
      }
      
      //  //book(_h["HIGGS_phi"]  ,20, 0, 2*M_PI);
      //  //book(_h["jet0_pT"], "jet0_pT", 10, 0.0, 800.0);
      //  //book(_h["jet1_pT"], "jet1_pT", 10, 0.0, 400.0);
      
      
      //book STXS specific hist
      //book(_h["counter_QQ2HQQ_FWDH"], "counter_QQ2HQQ_FWDH", 1,0,1);
      book(_h["counter_"+QQ2HQQ_0J], "counter_"+QQ2HQQ_0J, 1,0,1);
      book(_h["counter_"+QQ2HQQ_1J], "counter_"+QQ2HQQ_1J, 1,0,1);
      book(_h["stxsallcats"], "stxsallcats",40,0,40);
      for (auto STXScategory : QQ2HQQ_GE2J_categories){
	book(_h["counter_"+STXScategory], "counter_"+STXScategory, 1,0,1);
	book(_h["sign_DPhijj_"+STXScategory], "sign_DPhijj"+STXScategory, {-M_PI,-0.5*M_PI,0,0.5*M_PI,M_PI});
        }
    }

    
    /// @brief Checks whether the input particle has a child with a given PDGID
    bool hasChild(const GenParticle *ptcl, int pdgID) {
      for (auto child : Particle(*ptcl).children())
        if (child.pid() == pdgID) return true;
      return false;
    }

  std::string replaceAll(std::string str = "",std::string subStr = "QQ2HQQ_GE2J_MJJ_60_120_PTHJJ_GT25"){
      size_t pos = 0;
      while ((pos = str.find(subStr, pos)) != std::string::npos) {
          str.replace(pos, subStr.length(), "");
          pos += subStr.length();
      }
      return str;
  }

  //get a map of Mjj PTH and pTHJJ cuts e.g.: QQ2HQQ_GE2J_MJJ_60_120_PTHJJ_GT25 >> { "MJJ" => { 60.000000, 120.00000 }, "PTHJJ" => { 25.000000 } 
  map<string, vector<double>> get_VBF2J_STXScut(std::string STXS_name = "QQ2HQQ_GE2J_MJJ_60_120_PTHJJ_GT25" ){
    std::string input0 = replaceAll(STXS_name,"QQ2HQQ_GE2J_");
    std::string input = replaceAll(input0,"GT");
    input+="_";
    std::regex patternRegex("([A-Za-z]+)_((?:\\d+_)*)");  // Regular expression to match the desired pattern
    std::smatch matches;  // Object to store the matching results

    std::vector<std::pair<std::string, std::vector<double>>> extractedData;  // Vector to store the extracted data

    // Iterate over the string and find all matches
    std::string::const_iterator searchStart(input.cbegin());
    while (std::regex_search(searchStart, input.cend(), matches, patternRegex)) {
        // Extract the string
        std::string str = matches[1];

        // Extract the numbers from the double array
        std::string numbersStr = matches[2];
        std::vector<double> numbers;
        std::regex numbersRegex("\\d+");
        std::sregex_iterator numbersIterator(numbersStr.begin(), numbersStr.end(), numbersRegex);
        std::sregex_iterator numbersEnd;
        for (; numbersIterator != numbersEnd; ++numbersIterator) {
            numbers.push_back(std::stod(numbersIterator->str()));
        }

        // Store the string and double array in the vector
        extractedData.emplace_back(str, numbers);

        // Update the search start position
        searchStart = matches.suffix().first;
    }
    map<string, vector<double>> STXScuts;
    // Print the extracted data
    for (const auto& data : extractedData) {

        auto cuts = data.second;
        STXScuts[data.first]=cuts;
        //std::cout << "String: " << data.first << ", Double Array: { ";
        //for (const auto& number : data.second) {
            //std::cout << number << " ";
        //}
        //std::cout << "}" << std::endl;
    }
    return STXScuts;
  }

    /// Perform the per-event analysis
    void analyze(const Event& event) {

      /////////////////////////////////////////////
      //////////////Get Particles///////////////////
      /////////////////////////////////////////////
      std::map<string,bool> passSelection;
      //bool passSTXS=true;
      //bool passfiducial30=true;
      //bool passfiducial20=true;
      passSelection["STXS"]=true;
      passSelection["fiducial20"]=true;
      passSelection["fiducial30"]=true;
      
      // Get the final state particles ordered by pT
      const Particles& fs = apply<FinalState>(event, "FS").particlesByPt();
      
      // Find a stable Higgs (mandatory)
      const auto higgsiter = std::find_if(fs.begin(), fs.end(), [](const Particle& p){ return p.pid() == PID::HIGGSBOSON; });
      if (higgsiter == fs.end()) {
       MSG_WARNING("FATAL: No stable Higgs found in event record.\n");
       vetoEvent;}
      const Particle& higgs = *higgsiter;

      /*//find higgs using method from https://gitlab.cern.ch/mbasso/higgs-smeft-generation/-/blob/reweight/share/plugins/stxs/HiggsTemplateCrossSections.cc#L169
      Particle higgs;
      int Nhiggs = 0;
      for (const GenParticle *ptcl : Rivet::HepMCUtils::particles(event.genEvent())) {
	if (!ptcl) continue;
          // a) Reject all non-Higgs particles
          if (!PID::isHiggs(ptcl->pdg_id())) continue;
          // b) select only the final Higgs boson copy, prior to decay
          if (ptcl->end_vertex() && !hasChild(ptcl, PID::HIGGS)) {
            //std::cout<<"Found a higgggggggs"<<std::endl;
            higgs = Particle(ptcl);
            ++Nhiggs;
          }
	  }*/

    



      // Loop over final state particles and fill various particle vectors
      Particles leptons, photons, jet_ptcls;
      for ( const Particle& ptcl : fs ) {
        // Do not include the Higgs in jet finding!
        if ( ptcl.pid() == PID::HIGGSBOSON ) continue;
        // Neutrinos not from hadronisation
        if ( ptcl.isNeutrino() && !ptcl.fromHadron() ) continue;
        // Electrons and muons not from hadronisation
        if ( ( ptcl.abspid() == PID::ELECTRON || ptcl.abspid() == PID::MUON ) && !ptcl.fromHadron() ) {
          leptons.push_back(ptcl);
          continue;
        }
        // Photons not from hadronisation
        if ( ptcl.abspid() == PID::PHOTON && !ptcl.fromHadron() ) {
          photons.push_back(ptcl);
          continue;
        }
        // Add particle to jet inputs
        jet_ptcls.push_back(ptcl);
      }

      /*
      // Match FS photons to leptons within cone R=0.1
      // If they are not 'dressing' photons, add to jet particle vector
      for ( const Particle& ph : photons ) {
        bool fsr_photon = false;
        for ( const Particle& lep : leptons ) {
          if ( deltaR(ph.momentum(),lep.momentum()) < 0.1 ){
            fsr_photon=true;
            continue;
          }
        }
        if ( !fsr_photon ) jet_ptcls.push_back(ph);
      }
      */

      // Let's build the jets! By hand...
      const PseudoJets pjs_in = mkPseudoJets(jet_ptcls);
      const fastjet::JetDefinition jdef(fastjet::antikt_algorithm, 0.4);
      const Jets alljets = mkJets(fastjet::ClusterSequence(pjs_in, jdef).inclusive_jets());

      std::map<string, double> ptj_;
      std::map<string, Jets> jetmap;
     
      
	for(std::vector<string>::iterator it=selectionSetup.begin(); it!=selectionSetup.end(); it++){
	  if ((*it)=="fiducial30" || (*it)=="STXS") ptj_[(*it)] = 30;
	  if ((*it)=="fiducial20") ptj_[(*it)] = 20;
	  double ptj=ptj_[(*it)];
	  const Jets jets = sortByPt(filterBy(alljets, Cuts::pT > ptj*GeV && Cuts::absrap < 4.7));
	  jetmap[(*it)]=jets;
	}
      

      /////////////////////////////////////////////
      ////////////////Observables///////////////////
      /////////////////////////////////////////////
      std::map<string,double> dphijjmap;
      std::map<string,double> mjjmap;
      std::map<string,double> ptHjjmap;
      std::map<string,double> dyjjmap;


      for(std::vector<string>::iterator it=selectionSetup.begin(); it!=selectionSetup.end(); it++){
      //signed DPhijj definition
	double dphijj=-10.;
	double mjj = -10.;
	double ptHjj = -10.;
	double dyjj = -10.;
	size_t njets=jetmap[(*it)].size();
	const Jets jets = jetmap[(*it)];
	if (njets>1){
	  if (jets[0].rapidity()>jets[1].rapidity()) {
	    dphijj = remainder( jets[0].phi() - jets[1].phi(), 2*M_PI );
	  }
	  else{
	    dphijj = remainder( jets[1].phi() - jets[0].phi(), 2*M_PI );
	  }
	  
	  mjj = (jets[0].mom() + jets[1].mom()).mass();
	  dyjj = fabs(jets[0].rapidity() - jets[1].rapidity());
	  ptHjj = (jets[0].mom() + jets[1].mom() + higgs.mom()).pT();
	}
	
	dphijjmap[(*it)]=dphijj;
	mjjmap[(*it)]=mjj;
	ptHjjmap[(*it)]=ptHjj;
	dyjjmap[(*it)]=dyjj;
      }


      /////////////////////////////////////////////
      ////////////////Selections///////////////////
      // STXS block for STXS selections and hist
      // fiducial block for fiducial selections
      /////////////////////////////////////////////


      //***********************STXS selections block*********************************//
      //STXS selction: Higgs: |y_H| <2.5; Jets: anti-kT with R=0.4 and pTj > 30 GeV
      //if (selectionSetup=="STXS") {
      //if (fabs(higgs.rapidity())>2.5) vetoEvent;
      if (fabs(higgs.rapidity())>2.5) passSelection["STXS"]=false; // THIS MUST BE FALSE
      if(passSelection["STXS"]==true){
	// get each STXS category and fill counter and kinematic distribution per STXS 
	if (jetmap["STXS"].size()==0){
	  _h["counter_"+QQ2HQQ_0J]->fill(0.5);
	  _h["stxsallcats"]->fill(0.5);
	}
	else if (jetmap["STXS"].size()==1){
	  _h["counter_"+QQ2HQQ_1J]->fill(0.5);
	  _h["stxsallcats"]->fill(1.5);
	}
	else if (jetmap["STXS"].size()>1){
	  for(std::vector<string>::iterator it=QQ2HQQ_GE2J_categories.begin();
	      it!=QQ2HQQ_GE2J_categories.end(); it++){
	    string STXScategory=(*it);
	    map<string, vector<double>> STXScuts = get_VBF2J_STXScut(STXScategory);
	    
	    //check if event belong to this STXS bin
	    bool eventInThisSTXS = true;
	    string cut_variable;
	    double event_var_value=-1;
	    for (const auto& cut : STXScuts) {
	      event_var_value=-1;
	      cut_variable = cut.first;
	      if (cut_variable=="MJJ") event_var_value = mjjmap["STXS"];
	      if (cut_variable=="PTHJJ")event_var_value = ptHjjmap["STXS"];
	      if (cut_variable=="PTH")event_var_value = higgs.pT();
	      if (event_var_value<cut.second[0]) {
		eventInThisSTXS = false;
		break;
	    }
              if (cut.second.size()>1) {
                if (event_var_value>cut.second[1]) {
                  eventInThisSTXS = false;
                  break;
                }
              }
	  }
	  if (eventInThisSTXS) {
	    _h["counter_"+STXScategory]->fill(0.5);
	    _h["sign_DPhijj_"+STXScategory]->fill(dphijjmap["STXS"]);
	    int index = it - QQ2HQQ_GE2J_categories.begin();
	    _h["stxsallcats"]->fill(index+-2.5);
	    
	  }
          }
	}
      }
      //}
      
      //***********************end STXS selections block*********************************//


      //***********************fiducial selections block*********************************//
      

      for(std::vector<string>::iterator it=selectionSetup.begin(); it!=selectionSetup.end(); it++){
	if((*it).compare("STXS")==0) continue; 
	Jets jets  = jetmap[(*it)];
	for (auto jet : jets){
	  if (fabs(jet.rapidity())>4.7) passSelection[(*it)]=false;
	}
	if (mjjmap[(*it)]<300*GeV) passSelection[(*it)]=false;
	if (dyjjmap[(*it)]<2 ) passSelection[(*it)]=false;
      }
      
      
      
      //***********************end fiducial selections block****************************//


      /////////////////////////////////////////////
      ////////////Fill Common Histograms///////////
      /////////////////////////////////////////////

      for(std::vector<string>::iterator it=selectionSetup.begin(); it!=selectionSetup.end(); it++){
	if(passSelection[*it]){
	  _h["xsec"+(*it)]->fill(0.5);
	  _h["pT_H"+(*it)]->fill(higgs.pT());
	  _h["Njets"+(*it)]->fill(jetmap[(*it)].size() + 0.5); // accounts for HEPData offset
	  _h["Mjj"+(*it)]->fill(mjjmap[*it]); 
	  _h["sign_DPhijj"+(*it)]->fill(dphijjmap[*it]);
	  _h["DYjj"+(*it)]->fill(dyjjmap[*it]);
	  _h["pTHjj"+(*it)]->fill(ptHjjmap[*it]);
	  _h2D["Mjj__signDPhijj"+(*it)]->fill(mjjmap[*it],dphijjmap[*it]);
	  _h2D["Mjj__pTH"+(*it)]->fill(mjjmap[*it],higgs.pT());
	  _h2D["DYjj__pTH"+(*it)]->fill(dyjjmap[*it],higgs.pT());
	  _h2D["Njet__pTH"+(*it)]->fill(jetmap[*it].size()+0.5,higgs.pT());
	}
      }
      //_h["jet0_pT"]->fill(njets>0 ? 0 : jets[0].pT());
      //_h["jet1_pT"]->fill(njets>1 ? 0 : jets[1].pT());
      //_h["Mjj"]->fill(njets>1 ? 0 : (jets[0].mom() + jets[1].mom()).mass());
      //_h["DYjj"]->fill(njets>1 ? 0 : fabs(jets[0].rapidity() - jets[1].rapidity()));
      //_h["signedDPhijj"]->fill(dphijj);
      
    }

    
    /// Normalise histograms etc., after the run
    void finalize() {
      const double sf = crossSectionPerEvent();
      //const double sf = crossSection() / femtobarn  / sumOfWeights();
      //std::cout<<"cross section"<<crossSection()<<std::endl;
      //std::cout<<"sum of weight"<<sumOfWeights()<<std::endl;
      //std::cout<<sf<<std::endl;
      for (auto hist : _h) {
       scale(hist.second, sf*1000);
      }
    }


  private:
  

    map<string, Histo1DPtr> _h;
    map<string, Histo2DPtr> _h2D;
    vector<string> QQ2HQQ_GE2J_categories = {"QQ2HQQ_GE2J_MJJ_0_60_PTHJJ_0_25",
                                     "QQ2HQQ_GE2J_MJJ_60_120_PTHJJ_0_25",
                                     "QQ2HQQ_GE2J_MJJ_120_350_PTHJJ_0_25",
                                     "QQ2HQQ_GE2J_MJJ_0_60_PTHJJ_GT25",
                                     "QQ2HQQ_GE2J_MJJ_60_120_PTHJJ_GT25",
                                     "QQ2HQQ_GE2J_MJJ_120_350_PTHJJ_GT25",

                                     "QQ2HQQ_GE2J_MJJ_350_700_PTH_0_200_PTHJJ_0_25",
                                     "QQ2HQQ_GE2J_MJJ_350_700_PTH_0_200_PTHJJ_GT25",
                                     "QQ2HQQ_GE2J_MJJ_700_1000_PTH_0_200_PTHJJ_0_25",
                                     "QQ2HQQ_GE2J_MJJ_700_1000_PTH_0_200_PTHJJ_GT25",
                                     "QQ2HQQ_GE2J_MJJ_1000_1500_PTH_0_200_PTHJJ_0_25",
                                     "QQ2HQQ_GE2J_MJJ_1000_1500_PTH_0_200_PTHJJ_GT25",
                                     "QQ2HQQ_GE2J_MJJ_GT1500_PTH_0_200_PTHJJ_0_25",
                                     "QQ2HQQ_GE2J_MJJ_GT1500_PTH_0_200_PTHJJ_GT25",

                                     "QQ2HQQ_GE2J_MJJ_350_700_PTH_GT200_PTHJJ_0_25",
                                     "QQ2HQQ_GE2J_MJJ_350_700_PTH_GT200_PTHJJ_GT25",
                                     "QQ2HQQ_GE2J_MJJ_700_1000_PTH_GT200_PTHJJ_0_25",
                                     "QQ2HQQ_GE2J_MJJ_700_1000_PTH_GT200_PTHJJ_GT25",
                                     "QQ2HQQ_GE2J_MJJ_1000_1500_PTH_GT200_PTHJJ_0_25",
                                     "QQ2HQQ_GE2J_MJJ_1000_1500_PTH_GT200_PTHJJ_GT25",
                                     "QQ2HQQ_GE2J_MJJ_GT1500_PTH_GT200_PTHJJ_0_25",
                                     "QQ2HQQ_GE2J_MJJ_GT1500_PTH_GT200_PTHJJ_GT25"};
    
    string QQ2HQQ_FWDH = "QQ2HQQ_FWDH";
    string QQ2HQQ_0J = "QQ2HQQ_0J";
    string QQ2HQQ_1J = "QQ2HQQ_1J";
    
    vector <string> selectionSetup={"fiducial20","fiducial30","STXS"};
    
    //string selectionSetup = "fiducial20";
    //string selectionSetup = "fiducial30";
    //string selectionSetup = "STXS";
  };


  // The hook for the plugin system
  //RIVET_DECLARE_PLUGIN(MY_ANALYSIS);
  DECLARE_RIVET_PLUGIN(MY_ANALYSIS);

}