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

//using HepMC::GenParticle;


namespace Rivet {
  // @authors Silvia Ferrario Ravasio <silvia.ferrario.ravasio@cern.ch>
  //          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);

    /// ----------------------------------------------------------------------------
    /// --- Functions to get the histogram names
    /// ----------------------------------------------------------------------------
    // for the STX selection
    string select_stxs_bin(double invmjj, double pth, double pthjj, double delphijj){
      int invmjjbin,pthbin,pthjjbin,delphijjbin;
      delphijjbin = (delphijj<M_PI/2)  ? 0:1;
      pthbin      = (pth     <200.  )  ? 0:1;
      pthjjbin    = (pthjj   <25.   )  ? 0:1;
      if(invmjj>350e0 && invmjj<700e0) {
	invmjjbin = 0;
      }else if(invmjj>700e0 && invmjj<1000e0) {
	invmjjbin = 1;
      }else if(invmjj>1000e0 && invmjj<1500e0) {
	invmjjbin = 2;
      }else if(invmjj>1500e0) {
	invmjjbin = 3;
      }else{
        invmjjbin = -1;
	assert(false && "Could not locate mjj bin: This should not happen \n");
      }
      string histname = "STXS" + stxsmjjstr[invmjjbin] + stxspthstr[pthbin]
	+ stxspthjjstr[pthjjbin] + stxsdelphijjstr[delphijjbin];

      return histname;
    }
    // for our fiducial analysis
    vector<string> select_histo_bin(double pth, double delphijj, double delyjj, int njets, string cutlabel){
      int delphijjbin, pthbin, delyjjbin, njetsbin,ptjbin;
      //here the counting starts from 1 
      if(pth < 80.){
	pthbin = 1;
      }else if(pth < 120.){
	pthbin = 2.;
      }else if(pth < 260.){
	pthbin = 3.;
      }else if(pth < 500.){
	pthbin = 4;
      }else if(pth < 850.){
	pthbin = 5;
      }else{
	pthbin = 6;
      }
      delphijjbin = static_cast<int>(delphijj/(0.25*M_PI)) + 1;
      delyjjbin   =  (delyjj<4) ? 1: std::min(static_cast<int>(delyjj) -4 + 2, 5);
      njetsbin    = std::min(njets - 2 + 1, 4);
      if(cutlabel == "fiducial20"){
	ptjbin =1;
      }else if (cutlabel == "fiducial30"){
	ptjbin = 2;
      }else{
        ptjbin = -1;
	assert(false && "Could not locate ptjbin bin: This should not happen \n");
      }
      vector<string> hnames = {"HISTO"+mjjstr[0]+pthstr[pthbin]+ptjstr[ptjbin],
			       "HISTO"+mjjstr[0]+delphijjstr[delphijjbin]+ptjstr[ptjbin],
			       "HISTO"+mjjstr[0]+delyjjstr[delyjjbin]+ptjstr[ptjbin],
			       "HISTO"+pthstr[0]+delyjjstr[delyjjbin]+ptjstr[ptjbin],
			       "HISTO"+pthstr[0]+njetstr[njetsbin]+ptjstr[ptjbin]};
      return hnames;
    }

    
    /// ----------------------------------------------------------------------------
    /// --- init: book histograms and define particles projectors
    /// ----------------------------------------------------------------------------
    void init() {

      // All final state particles
      const FinalState fs;
      declare(fs, "FS");
      vector<double> mjjbins      = {300., 500., 700., 900., 1100, 1e50};
      vector<double> ptHbins      = {0., 80., 120, 260., 500., 850., 1e50};
      //vector<double> delyjjbins   = {2., 4., 5., 6., 7., 1e50};
      //vector<double> delphijjbins = {0., M_PI/4., M_PI/2., 3.*M_PI/4., M_PI};
      //vector<double> njetbins     = {2., 3., 4., 5., 6.};

      // Inclusive histos      
      book(_h["sig incl cuts (ptj > 20 GeV)"], "sig incl cuts (ptj > 20 GeV)", 1, 0, 1);
      book(_h["sig incl cuts (ptj > 30 GeV)"], "sig incl cuts (ptj > 30 GeV)", 1, 0, 1);

      // STXS bins
      for(int i = 0; i < stxsmjjstr.size(); i++){
	for(int j = 0; j < stxspthstr.size(); j++){
	  for(int k = 0; k< stxspthjjstr.size(); k++){
	    for (int l = 0; l<stxsdelphijjstr.size(); l++){
		string histostr = "STXS"+stxsmjjstr[i]+stxspthstr[j]+stxspthjjstr[k]+stxsdelphijjstr[l];
		book(_h[histostr], histostr, 1, 0, 1);
	      }
	  }
	}
      }

      //Our own bins. The two ptjstr values in each histograms correspond to ptj>20 and ptj>30
      //Mjj 
      int i=0;
      //.. vs pth, dphijj, dyjj
      for( auto stringvector : {pthstr, delphijjstr, delyjjstr}){
	for(int j=1; j<stringvector.size(); j++){
	  for(int k=1; k<ptjstr.size(); k++){
	    string histostr = "HISTO"+mjjstr[i]+stringvector[j]+ptjstr[k];
	    book(_h[histostr], histostr, mjjbins);
	  }
	}
      }
      //ptH
      i=0;
      //..vs dyjj, Njj
      for( auto stringvector : {delyjjstr, njetstr}){
	for(int j=1; j<stringvector.size(); j++){
	  for(int k=1; k<ptjstr.size(); k++){
	    string histostr = "HISTO"+pthstr[i]+stringvector[j]+ptjstr[k];
	    book(_h[histostr], histostr, ptHbins);
	  }
	}
      }      
      //Additional histos
      for( auto ptstr : {"ptj20", "ptj30"}){
	string histostr = "ptHjj-"+std::string(ptstr);
	book(_h[histostr], histostr, 50, 0e0, 1000e0);
	histostr+="-overflow";
	book(_h[histostr], histostr, 1., 0e0, 1e0);	
      }

      // Extra histograms for comparisons with Mathieu. These use the STXS cuts.
      book(_h["ptj1-STXS"], "ptj1-STXS", 100, 0e0, 2000e0);
      book(_h["ptj2-STXS"], "ptj2-STXS", 100, 0e0, 2000e0);
      book(_h["mjj-STXS"], "mjj-STXS", 100, 0e0, 5000e0);
      book(_h["yjj-STXS"], "yjj-STXS", 36, -4.5e0, 4.5e0);
      book(_h["ptH-STXS"], "ptH-STXS", 100, 0e0, 2000e0);
      book(_h["ptHjj-STXS"], "ptHjj-STXS", 50, 0e0, 1000e0);
      book(_h["yj1-STXS"], "yj1-STXS", 36, -4.5e0, 4.5e0);
      book(_h["yj2-STXS"], "yj2-STXS", 36, -4.5e0, 4.5e0);
      book(_h["yH-STXS"], "yH-STXS", 36, -4.5e0, 4.5e0);

      
    }

    /// ----------------------------------------------------------------------------
    /// --- Analyse the event
    /// ----------------------------------------------------------------------------
    void analyze(const Event& event) {

      std::map<string,bool> passSelection = {{"STXS", true}, {"fiducial20", true}, {"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;
      double yH  = higgs.rap();
      double ptH = higgs.pt();
      
      if(print_debug){
	std::cout<<"***************************************************\n";
	std::cout<<"Event \n";
	std::cout<<"Higgs: \n"<<higgs<<"\n";
	std::cout<<"Higgs rap: "<<yH<<"\n";
      }
      
      // 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);
      }

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

      std::map<string, double> mjjmap;
      std::map<string, double> dyjjmap;
      std::map<string, double> dphijjmap;
      std::map<string, double> ptHjjmap;
      std::map<string, int>    njetsmap;
      std::map<string, Jets>   jetsmap;
      
     
      
      for(std::vector<string>::iterator it=selectionSetup.begin(); it!=selectionSetup.end(); it++){
	//First check if the Higgs passes the cuts for a given setup
	if(std::abs(yH) > absyHmax[(*it)]){
	  passSelection[(*it)]=false;
	  continue;
	}
	
	const Jets jets = sortByPt(filterBy(alljets, Cuts::pT > ptjmin[(*it)]*GeV
					    && Cuts::absrap < absyjmax[(*it)]));
	if(jets.size()<2){
	  passSelection[(*it)]=false;
	  continue;
	}
	njetsmap[(*it)]=jets.size();
	jetsmap[(*it)]=jets;
      
	double mjj =  (jets[0].mom() + jets[1].mom()).mass();
	if(mjj<mjjmin[(*it)]){
	  passSelection[(*it)]=false;
	  continue;
	}
	mjjmap[(*it)]=mjj;
      
	double dyjj = std::abs(jets[0].rap()-jets[1].rap());
	if(dyjj<dyjjmin[(*it)]){
	  passSelection[(*it)]=false;
	  continue;
	}
	dyjjmap[(*it)]=dyjj;

	//Calculate dphijj using the same convention as in the powheg analysis
	double dphijj = std::abs(jets[0].phi()-jets[1].phi());
	dphijj = min(dphijj, std::abs(2.*M_PI-dphijj));
	dphijjmap[(*it)]=dphijj;

	ptHjjmap[(*it)]=(jets[0].mom()+jets[1].mom()+higgs.mom()).pt();
      }

      if(print_debug){
	for(std::vector<string>::iterator it=selectionSetup.begin(); it!=selectionSetup.end(); it++){
	  std::cout<<(*it)<<" "<< passSelection[(*it)]<<"\n";
	}
	std::cout<<"***************************************************\n\n";
      }

      
      if(passSelection["STXS"]){
	string sxtshisto =  select_stxs_bin(mjjmap["STXS"], ptH, ptHjjmap["STXS"], dphijjmap["STXS"]);
	_h[sxtshisto]->fill(0.5);
	_h["ptj1-STXS"]->fill(jetsmap["STXS"][0].pt());
	_h["ptj2-STXS"]->fill(jetsmap["STXS"][1].pt());
	_h["mjj-STXS"]->fill(mjjmap["STXS"]);
	_h["yjj-STXS"]->fill(jetsmap["STXS"][0].rap()-jetsmap["STXS"][1].rap()); //Signed rapidity
	_h["ptHjj-STXS"]->fill(ptHjjmap["STXS"]);
	_h["yj1-STXS"]->fill(jetsmap["STXS"][0].rap());
	_h["yj2-STXS"]->fill(jetsmap["STXS"][1].rap());
	_h["yH-STXS"]->fill(yH);
	_h["ptH-STXS"]->fill(ptH);
      }


      for(auto ptc : {"20", "30"}){
	string cutlabel = "fiducial"+std::string(ptc);
	if(not passSelection[cutlabel]) continue;
	
	_h["sig incl cuts (ptj > "+std::string(ptc)+" GeV)"]->fill(0.5);
 
	vector<string> myhistos = select_histo_bin(ptH, dphijjmap[cutlabel], dyjjmap[cutlabel], njetsmap[cutlabel],
						   cutlabel);
	_h[myhistos[0]]->fill(mjjmap[cutlabel]);
	_h[myhistos[1]]->fill(mjjmap[cutlabel]);
	_h[myhistos[2]]->fill(mjjmap[cutlabel]);
	_h[myhistos[3]]->fill(ptH);
	_h[myhistos[4]]->fill(ptH);
	
	
	if(ptHjjmap[cutlabel]<1000.){
	  _h["ptHjj-ptj"+std::string(ptc)]->fill(ptHjjmap[cutlabel]);
	}else{
	  _h["ptHjj-ptj"+std::string(ptc)+"-overflow"]->fill(0.5);
	}
	
      }
    }

    

    /// ----------------------------------------------------------------------------
    /// --- Finalise: normalise the histograms
    /// ----------------------------------------------------------------------------
    void finalize() {
      const double sf = crossSectionPerEvent();
      for (auto hist : _h) {
       scale(hist.second, sf*1000);
      }
    }


  private:
    map<string, Histo1DPtr> _h;

    vector <string> selectionSetup = {"STXS", "fiducial20", "fiducial30"};

    //smallest pT on jet
    std::map<string, double> ptjmin  = {{"STXS",       30.},
					{"fiducial20", 20.},
					{"fiducial30", 30.}};
    //largest |y| on jet
    std::map<string, double> absyjmax = {{"STXS",       47.},
					 {"fiducial20", 4.7},
					 {"fiducial30", 4.7}};  
    //radius in the jet algorithm
    std::map<string, double> Rjet     = {{"STXS",       0.4},
					 {"fiducial20", 0.4},
					 {"fiducial30", 0.4}};
    //min mass of the leading jets pair
    std::map<string, double> mjjmin   = {{"STXS",       350.},
					 {"fiducial20", 300.},
					 {"fiducial30", 300.}} ;
    //min delta y of the leading jets pair
    std::map<string, double> dyjjmin  = {{"STXS",       0.},
					 {"fiducial20", 2.},
					 {"fiducial30", 2.}};

    //largest |y| on the Higgs
    std::map<string, double> absyHmax = {{"STXS",       2.5},
					 {"fiducial20", 5e50},
					 {"fiducial30", 5e50}};
    
    

    //Same names as "setup_bin_names()" in powheg analysis
    const vector<string> stxsmjjstr      = {"-MJJ-350-700",  "-MJJ-700-1000", "-MJJ-1000-1500", "-MJJ-1500-INFTY"};
    const vector<string> stxspthstr      = {"-PTH-0-200", "-PTH-200-INFTY"};
    const vector<string> stxspthjjstr    = {"-PTHJJ-0-25", "-PTHJJ-25-INFTY"};
    const vector<string> stxsdelphijjstr = {"-DPHIJJ-0-PIov2", "-DPHIJJ-PIov2-PI"};
    const vector<string> mjjstr          = {"-MJJ", "-MJJ-300-500", "-MJJ-500-700", "-MJJ-700-900", "-MJJ-900-1100", "-MJJ-1100-INFTY"};
    const vector<string> pthstr          = {"-PTH", "-PTH-0-80", "-PTH-80-120", "-PTH-120-260", "-PTH-260-500", "-PTH-500-850", "-PTH-850-INFTY"};
    const vector<string> delphijjstr     = {"-DPHIJJ", "-DPHIJJ-0-PIov4", "-DPHIJJ-PIov4-PIov2", "-DPHIJJ-PIov2-3PIov4", "-DPHIJJ-3PIov4-PI"};
    const vector<string> delyjjstr       = {"-DYJJ", "-DYJJ-2-4", "-DYJJ-4-5", "-DYJJ-5-6", "-DYJJ-6-7", "-DYJJ-7-INFTY"};
    const vector<string> njetstr         = {"-NJETS", "-NJETS-2", "-NJETS-3", "-NJETS-4", "-NJETS-5"};
    const vector<string> ptjstr          = {"-PTJ", "-PTJ-20", "-PTJ-30"};




    const bool print_debug = false;
    
  };


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

}