//!-----------------------------------
/*
  Authors: Raghav Kunnawalkam Elayavalli and Korinna Christine Zapp 
  July 5th 2017, 
  RIVET Analysis for Jet - Background subtraction in JEWEL
  To be run with jewel-v2.2.0, with including recoils, 
  dummy particles and the scattering centers (HepMC:StatusCode 3)  
*/
//!-----------------------------------

#include "Rivet/Analysis.hh"
#include "Rivet/Projections/FinalState.hh"
#include "Rivet/Projections/FastJets.hh"
#include "Rivet/Tools/Logging.hh"
#include "Rivet/Tools/ParticleIdUtils.hh"
#include "fastjet/AreaDefinition.hh"
#include "fastjet/ClusterSequence.hh"
#include "fastjet/ClusterSequenceArea.hh"
#include "HepMC/GenParticle.h"
#include "HepMC/GenEvent.h"
#include "HepMC/GenCrossSection.h"


namespace Rivet {


  class JEWEL_BKGSUBTRACTION : public Analysis {
  public:

    //! Constructor
    JEWEL_BKGSUBTRACTION(string name = "JEWEL_BKGSUBTRACTION")
      : Analysis(name)
    {
      setNeedsCrossSection(true);
    }

    //! Cells which make up the grid  
    struct candidate{
      //! these are the boundaries 
      double etaMin;
      double phiMin;
      double etaMax;
      double phiMax;
      //! Pseudojets will have the objects inside 
      PseudoJets objects;
      vector <int> jetID;
      //! Four momeentum to get useful informations like eta, phi, mass, pT etc... 
      FourMomentum candMom;
      FourMomentum bkgMom;
      double sumnegpT;
      double sumpT;
    };

    std::pair<int,int> FindCandidateBin(double eta, double phi){
      int etaloc = -9;
      int philoc = -9;      
      for(unsigned x = 0; x<_etabins.size()-1; ++x){
	for(unsigned y = 0; y<_phibins.size()-1; ++y){
	  if(eta >= _etabins[x] && eta < _etabins[x+1] &&
	     phi >= _phibins[y] && phi < _phibins[y+1]){
	    etaloc = x;
	    philoc = y;
	  }	  
	}
      }
      if(etaloc == -9 || philoc == -9)
	return std::pair<int,int>(-1,-1);
      else
	return std::pair<int,int>(etaloc,philoc);
    }

    //! 4-Momentum Subtraction Procedure
    PseudoJets do4MomSub(const Jets inputCollection,
			 vector <HepMC::GenParticle*> scatteringCenters,
			 bool doSubtraction){

      PseudoJets jets_bkgsub;
      //! loop over the input jet collection
      foreach (const PseudoJet jet, inputCollection){
	FourMomentum jt, bkg, jt_sub;
	//! constituents map	
	foreach (const fastjet::PseudoJet &c, jet.constituents()){
	  jt += FourMomentum(c.E(), c.px(), c.py(), c.pz());
	}
	    
	if(doSubtraction){
	  for(unsigned ip = 0; ip<scatteringCenters.size(); ++ip){
	    bool isinJet = false;	  
	    double _delRSmall = 100;
	    double _delRposition = 0;
	    int counter = 0;
	    foreach (const fastjet::PseudoJet &c, jet.constituents()){
	      const double _delRTruth = deltaR(scatteringCenters[ip]->momentum().eta(),
					       scatteringCenters[ip]->momentum().phi(),
					       c.eta(), c.phi_std());
	      if(_delRTruth < 1e-5){
		_delRSmall = _delRTruth;
		_delRposition = counter;
		isinJet = true;
	      }
	      counter++;
	    }// jet constituents loop
	    // sum the background contribution
	    if(isinJet){
	      bkg += FourMomentum(scatteringCenters[ip]->momentum().e(),
				  scatteringCenters[ip]->momentum().px(),
				  scatteringCenters[ip]->momentum().py(),
				  scatteringCenters[ip]->momentum().pz());
	    }
	  }//! scattering center loop
	      
	  jt_sub = jt - bkg;
	}else
	  jt_sub = jt;
	    
	if(jt_sub.pT() <= _pTCut) continue;
	jets_bkgsub.push_back(PseudoJet(jt_sub.px(), jt_sub.py(), jt_sub.pz(), jt_sub.E()));	  	
      }

      if(jets_bkgsub.size()==0){
	if(verbose) std::cout<<"no 4MomSub jets in the event with given kinematic cuts"<<std::endl;
      }
      
      return jets_bkgsub;

    }

    //! Grid Based Subtraction - 1 Procedure (Jet Clustering before Discretization)
    PseudoJets doGridSub1(const Jets inputCollection,
			  vector <HepMC::GenParticle*> scatteringCenters,
			  bool doSubtraction){

      //! initialize the grid
      vector <vector<candidate> > grid;
      for(int x = 0; x<_Nbounds_eta; ++x){
	vector <candidate> xgrid;
	for(int y = 0; y<_Nbounds_phi; ++y){
	  candidate test;
	  test.etaMin = _etabins[x];
	  test.phiMin = _phibins[y];
	  test.etaMax = _etabins[x+1];
	  test.phiMax = _phibins[y+1];
	  test.sumnegpT = 0.0;
	  xgrid.push_back(test);
	}
	grid.push_back(xgrid);
	xgrid.clear();
      }
      
      //! Now loop over all the jets and fill their constituents into the grid
      int jetcounter = 0;      
      foreach ( const PseudoJet jet, inputCollection ) {
	jetcounter++;
	//! loop over the candidates of the jet and include them in the grid
	foreach ( const PseudoJet c, jet.constituents() ){
	  double ceta = c.eta();
	  double cphi = c.phi_std();
	  std::pair<int,int> candpos = FindCandidateBin(ceta, cphi);
	  if(candpos.first!=-1 && candpos.second!=-1){
	    grid[candpos.first][candpos.second].objects.push_back(c);
	    grid[candpos.first][candpos.second].candMom+=FourMomentum(c.E(), c.px(), c.py(), c.pz());
	    grid[candpos.first][candpos.second].jetID.push_back(jetcounter);
	    grid[candpos.first][candpos.second].sumpT+=c.pt();
	  }
	}
      }// jet loop
	  
      //! Now loop over the scattering centers that are near jets and add that as BKG
      if(doSubtraction){
	for(unsigned ip = 0; ip<scatteringCenters.size(); ++ip){
	  double px,py,pz,E;
	  px = scatteringCenters[ip]->momentum().px();
	  py = scatteringCenters[ip]->momentum().py();
	  pz = scatteringCenters[ip]->momentum().pz();
	  E  = scatteringCenters[ip]->momentum().e();
	  PseudoJet part(px,py,pz,E);
	  bool iJ = false;
	  foreach(const PseudoJet j, inputCollection){
	    double delR = deltaR(j.eta(), j.phi_std(),
				 part.eta(), part.phi_std());
	    if(delR < _jetR){
	      iJ = true;
	      break;
	    }
	  }
	  if(iJ){
	    std::pair<int,int> candpos = FindCandidateBin(part.eta(), part.phi_std());
	    if(candpos.first!=-1 && candpos.second!=-1){
	      grid[candpos.first][candpos.second].objects.push_back(part);
	      grid[candpos.first][candpos.second].bkgMom+=FourMomentum(E, px, py, pz);
	      grid[candpos.first][candpos.second].sumpT-=part.pt();
	    }
	  }
	}// scattering center loop
      }// dosub
	  
      //! build up the pseudojets to do the clustering
      PseudoJets pJet_sub;
      for(int x = 0; x<_Nbounds_eta; ++x){
	for(int y = 0; y<_Nbounds_phi; ++y){
	  FourMomentum jet_sub = grid[x][y].candMom;
	  if(doSubtraction)
	    jet_sub -= grid[x][y].bkgMom;
	  double px, py, pz, E;
	  if(doSubtraction && (grid[x][y].candMom.pT() - grid[x][y].bkgMom.pT() < 0.0)){
	    grid[x][y].sumnegpT += grid[x][y].bkgMom.pT() - grid[x][y].candMom.pT();
	    px = 0.0; py = 0.0; pz = 0.0; E = 0.0;
	  }else {
	    px = jet_sub.px();
	    py = jet_sub.py();
	    pz = jet_sub.pz();
	    E  = jet_sub.E();
	  }
	  PseudoJet part_sub(px,py,pz,E);
	  if(part_sub.pt() != 0 && part_sub.E() != 0)
	    pJet_sub.push_back(part_sub);
	}
      }//!grid loop
      
      fastjet::JetDefinition jet_def(fastjet::antikt_algorithm, _jetR);
      fastjet::ClusterSequence cs_sub(pJet_sub, jet_def);
      PseudoJets jets_bkgsub = sorted_by_pt(cs_sub.inclusive_jets(_pTCut));

      if(jets_bkgsub.size()==0){
	if(verbose) std::cout<<"no GridSub1 jets in the event with given kinematic cuts"<<std::endl;
      }

      return jets_bkgsub; 

    }

    //! Grid Based Subtraction - 2 Procedure (Discretization before jet clustering)
    PseudoJets doGridSub2(const Event& event,
			  const Jets inputCollection,
			  vector <HepMC::GenParticle*> scatteringCenters,
			  bool doSubtraction){

      
      //! initialize the grid
      vector <vector<candidate> > grid;
      for(int x = 0; x<_Nbounds_eta; ++x){
	vector <candidate> xgrid;
	for(int y = 0; y<_Nbounds_phi; ++y){
	  candidate test;
	  test.etaMin = _etabins[x];
	  test.phiMin = _phibins[y];
	  test.etaMax = _etabins[x+1];
	  test.phiMax = _phibins[y+1];
	  test.sumnegpT = 0.0;
	  xgrid.push_back(test);
	}
	grid.push_back(xgrid);
	xgrid.clear();
      }

      //! Now loop over all the final state particles and fill their constituents into the grid
      const ParticleVector& FS = applyProjection<FinalState>(event, "FS").particlesByPt();
      foreach ( const Particle& p, FS) {
	double peta = p.eta();
	double pphi = p.phi(MINUSPI_PLUSPI);
	double ppT  = p.pT();

	std::pair<int,int> candpos = FindCandidateBin(peta, pphi);
	if(candpos.first!=-1 && candpos.second!=-1){
	  grid[candpos.first][candpos.second].objects.push_back(PseudoJet(p.px(), p.py(), p.pz(), p.E()));
	  grid[candpos.first][candpos.second].candMom+=FourMomentum(p.E(), p.px(), p.py(), p.pz());
	  grid[candpos.first][candpos.second].sumpT+=ppT;
	}
      }

      //! This part of GridSub2 is same as GridSub1 
      //! Now loop over the scattering centers that are near jets and add that as BKG
      if(doSubtraction){
	for(unsigned ip = 0; ip<scatteringCenters.size(); ++ip){
	  double px,py,pz,E;
	  px = scatteringCenters[ip]->momentum().px();
	  py = scatteringCenters[ip]->momentum().py();
	  pz = scatteringCenters[ip]->momentum().pz();
	  E  = scatteringCenters[ip]->momentum().e();
	  PseudoJet part(px,py,pz,E);
	  bool iJ = false;
	  foreach(const PseudoJet j, inputCollection){
	    double delR = deltaR(j.eta(), j.phi_std(),
				 part.eta(), part.phi_std());
	    if(delR < _jetR){
	      iJ = true;
	      break;
	    }
	  }
	  if(iJ){
	    std::pair<int,int> candpos = FindCandidateBin(part.eta(), part.phi_std());
	    if(candpos.first!=-1 && candpos.second!=-1){
	      grid[candpos.first][candpos.second].objects.push_back(part);
	      grid[candpos.first][candpos.second].bkgMom+=FourMomentum(E, px, py, pz);
	      grid[candpos.first][candpos.second].sumpT-=part.pt();
	    }
	  }
	}// scattering center loop
      }// dosub
	  
      //! build up the pseudojets to do the clustering
      PseudoJets pJet_sub;
      for(int x = 0; x<_Nbounds_eta; ++x){
	for(int y = 0; y<_Nbounds_phi; ++y){
	  FourMomentum jet_sub = grid[x][y].candMom;
	  if(doSubtraction)
	    jet_sub -= grid[x][y].bkgMom;
	  double px, py, pz, E;
	  if(doSubtraction && (grid[x][y].candMom.pT() - grid[x][y].bkgMom.pT() < 0.0)){
	    grid[x][y].sumnegpT += grid[x][y].bkgMom.pT() - grid[x][y].candMom.pT();
	    px = 0.0; py = 0.0; pz = 0.0; E = 0.0;
	  }else {
	    px = jet_sub.px();
	    py = jet_sub.py();
	    pz = jet_sub.pz();
	    E  = jet_sub.E();
	  }
	  PseudoJet part_sub(px,py,pz,E);
	  if(part_sub.pt() != 0 && part_sub.E() != 0)
	    pJet_sub.push_back(part_sub);
	}
      }//!grid loop
      
      fastjet::JetDefinition jet_def(fastjet::antikt_algorithm, _jetR);
      fastjet::ClusterSequence cs_sub(pJet_sub, jet_def);
      PseudoJets jets_bkgsub = sorted_by_pt(cs_sub.inclusive_jets(_pTCut));

      if(jets_bkgsub.size()==0){
	if(verbose) std::cout<<"no GridSub2 jets in the event with given kinematic cuts"<<std::endl;
      }

      return jets_bkgsub; 
      
    }
    
    
    //! Book histograms and initialise projections before the run
    void init() {

      //! verbosity flag
      verbose = false;
      
      //! Jet Radius
      _jetR=0.4;
      //! Jet Kinematic cuts 
      _pTCut=100.0;

      //! Grid initialization  
      _etaMax=2.4;
      _phiMax=M_PI;
      _delRMin=0.05;

      //! cell boundaries 
      _Nbounds_eta = (int)2*_etaMax/(_delRMin);
      _Nbounds_phi = (int)2*_phiMax/(_delRMin);
      for(int i = 0; i<=_Nbounds_eta; ++i){
	double etax = -1*_etaMax + i*_delRMin;
	_etabins.push_back(etax);
      }
      for(int i = 0; i<=_Nbounds_phi; ++i){
	double phix = -1*_phiMax + i*_delRMin;
	_phibins.push_back(phix);	
      }
      
      if(verbose)
	std::cout<<"Grid we are using is ("<<_Nbounds_eta<<" x "<<_Nbounds_phi
		 <<")   eta: [-"<<_etaMax<<", "<<_etaMax
		 <<"]   phi: [-"<<_phiMax<<", "<<_phiMax<<"]"<<std::endl;      

      //! Initialize Histograms
      //! Example jet pT histograms for 3 jet collections 
      _h_JetpT_noSub    = bookHisto1D("JetpT_noSub", 30, 100, 400);
      _h_JetpT_4MomSub  = bookHisto1D("JetpT_4MomSub", 30, 100, 400);
      _h_JetpT_GridSub1 = bookHisto1D("JetpT_GridSub1", 30, 100, 400);
      _h_JetpT_GridSub2 = bookHisto1D("JetpT_GridSub2", 30, 100, 400);

      //! final state and jet projection
      FinalState fs(-3.0, 3.0, 0.*GeV);
      addProjection(fs, "FS");
      addProjection(FastJets(fs, FastJets::ANTIKT, _jetR), "Jets");      
      
    }

    //! Perform the per-event analysis
    void analyze(const Event& event) {
      const double weight = event.weight();

      //! **************************************
      //! Jet Collection from all particles in the event
      //! Used for w/o recoils and in vacuum 
      if(verbose) std::cout<<"Jet Collection built without subtraction"<<std::endl;
      const FastJets& AJets = applyProjection<FastJets>(event, "Jets");
      Cut cuts = Cuts::etaIn(-2.0, 2.0) & (Cuts::pT > _pTCut*GeV);
      const Jets jets_noSub = AJets.jetsByPt(cuts);
      foreach ( const PseudoJet jet, jets_noSub ) {
	_h_JetpT_noSub->fill(jet.pt(), weight);
      }

            
      //! Checking if additional scattering centers are available from HepMC file 
      //! get the scattered particles from the gen event:
      vector <HepMC::GenParticle*> pscat; 
      foreach (const HepMC::GenParticle* p, particles(event.genEvent())) {
	if(p->status() == 3)
	  pscat.push_back((HepMC::GenParticle*)p);
      }
	
      bool doSubtraction = true;
      if(pscat.size() == 0)
	doSubtraction = false;

      
      //! **************************************
      //! BKG-Sub Jet Collection from all particles w/ recoils inclding the dummy and Scattering centers 
      if(verbose) std::cout<<"Jet Collection w/ 4MomSub Subtraction"<<std::endl;
      PseudoJets jets_4MomSub = do4MomSub(jets_noSub, pscat, doSubtraction);
      foreach ( const PseudoJet jet, jets_4MomSub ) {
	_h_JetpT_4MomSub->fill(jet.pt(), weight);
      }

      //! **************************************
      //! BKG-Sub Jet Collection from all particles w/ recoils inclding the Scattering centers 
      if(verbose) std::cout<<"Jet Collection w/ GridSub1 Subtraction"<<std::endl;
      PseudoJets jets_GridSub1 = doGridSub1(jets_noSub, pscat, doSubtraction);
      foreach ( const PseudoJet jet, jets_GridSub1 ) {
	_h_JetpT_GridSub1->fill(jet.pt(), weight);
      }

      //! **************************************
      //! BKG-Sub Jet Collection from all particles w/ recoils inclding the Scattering centers 
      if(verbose) std::cout<<"Jet Collection w/ GridSub2 Subtraction"<<std::endl;
      PseudoJets jets_GridSub2 = doGridSub2(event, jets_noSub,  pscat, doSubtraction);
      foreach ( const PseudoJet jet, jets_GridSub2 ) {
	_h_JetpT_GridSub2->fill(jet.pt(), weight);
      }
      

    }


    //! Normalise histograms etc., after the run
    void finalize() {

      scale(_h_JetpT_noSub, crossSection()/picobarn/sumOfWeights());
      scale(_h_JetpT_4MomSub, crossSection()/picobarn/sumOfWeights());
      scale(_h_JetpT_GridSub1, crossSection()/picobarn/sumOfWeights());
      scale(_h_JetpT_GridSub2, crossSection()/picobarn/sumOfWeights());

    }

  protected:

    double _pTCut;
    double _jetR;
    bool verbose;
    vector<double> _etabins;
    vector<double> _phibins;
    int _Nbounds_eta;
    int _Nbounds_phi;
    double _delRMin;
    double _etaMax;
    double _phiMax;


  private:

    Histo1DPtr _h_JetpT_noSub;
    Histo1DPtr _h_JetpT_4MomSub;
    Histo1DPtr _h_JetpT_GridSub1;
    Histo1DPtr _h_JetpT_GridSub2;


  };



  //! The hook for the plugin system
  DECLARE_RIVET_PLUGIN(JEWEL_BKGSUBTRACTION);


}