// -*- C++ -*-
#include "Rivet/Analysis.hh"
#include "Rivet/RivetAIDA.hh"
#include "Rivet/Tools/Logging.hh"
#include "Rivet/Projections/FastJets.hh"
#include "Rivet/Tools/BinnedHistogram.hh"

#include "Rivet/Particle.hh"
#include "Rivet/Tools/ParticleIdUtils.hh"
#include "HepMC/GenEvent.h"


namespace Rivet {


  /// @brief ATLAS inclusive b-jet pT spectrum, di-jet Mass and di-jet chi
  class ATLAS_2011_I930220: public Analysis{
  public:

    ATLAS_2011_I930220() : Analysis("ATLAS_2011_I930220")
    {
      setBeams(PROTON,PROTON);
      setNeedsCrossSection(true);
    }


  public:

    void init() {

      FinalState fs(-3.5,3.5);
      addProjection(fs, "FinalState");

      FastJets fj(fs, FastJets::ANTIKT, 0.4);
      fj.useInvisibles();
      addProjection(fj, "Jets");

      double ybins[] = {0.0, 0.3, 0.8, 1.2, 2.1};
      
      for (size_t i = 0; i < 4; ++i) {
	_bjetpT_SV0.addHistogram(ybins[i], ybins[i+1], bookHistogram1D(i + 1, 1, 1) );
      }

      _bjetpT_SV0_All    = bookHistogram1D(5, 1, 1);
      _bjetpT_pTRel      = bookHistogram1D(6, 1, 1);
      _dijet_mass        = bookHistogram1D(7, 1, 1);
      _dijet_phi         = bookHistogram1D(8, 1, 1);
      _dijet_chi_110_370 = bookHistogram1D(9, 1, 1);
      _dijet_chi_370_850 = bookHistogram1D(10, 1, 1);

      ChiCounter1 = 0.0;
      ChiCounter2 = 0.0;
      PhiCounter  = 0.0;

    }      
    
    void analyze(const Event& evt) {

      const double weight = evt.weight();

      //Finding all b-hadrons within event
      std::vector<HepMC::GenParticle*> bHadrons;
      std::vector<HepMC::GenParticle*> allParticles = particles(evt.genEvent());
      for (unsigned int i = 0; i < allParticles.size(); i++)  {
	GenParticle* p = allParticles.at(i);
	if ( ! ( Rivet::PID::isHadron( p->pdg_id() ) && Rivet::PID::hasBottom( p->pdg_id() ) ) ) continue;
	if ( p->momentum().perp()*GeV < 5 ) continue;
	bHadrons.push_back(p);
      }

      FourMomentum LeadingJet;
      FourMomentum SubLeadingJet;
      int leadJet = 0;
      int subJet  = 0;       

      const Jets jets = applyProjection<JetAlg>(evt, "Jets").jetsByPt(15.0*GeV);
      foreach(const Jet& j, jets) {

	FourMomentum jet = j.momentum();
	double  y = fabs(jet.rapidity());
	double Pt = jet.pT()*GeV; 
	bool hasB = false;

	for (unsigned int i = 0; i < bHadrons.size(); i++) {
	  FourVector bHadron = bHadrons.at(i)->momentum();
	  double dR = deltaR(jet.eta(),jet.phi(),bHadron.eta(),bHadron.phi());
 	  if( dR < 0.3 ) {	    
	    hasB = true;
	    break;
 	  }
	}
	
	/*Event is suitable for the di-jet measurment if both the leading and sub-leading jet in the event
	  are b-tagged and have a pT > 40 GeV*/
	if ( y < 2.1 ) {
	  if ( leadJet && !subJet ) {
	    SubLeadingJet = jet;
	    subJet = (hasB && Pt > 40) ? 2 : 1;
	  }
	  if ( ! leadJet ) {
	    LeadingJet = jet;
	    leadJet = (hasB && Pt > 40) ? 2 : 1;
	  }	
	  if ( hasB ) {
	    _bjetpT_SV0.fill(y,Pt, weight);
	    _bjetpT_SV0_All->fill(Pt, weight);
	    _bjetpT_pTRel->fill(Pt, weight);	    
	  }
	}
      }

      if ( (leadJet == 2) && (subJet == 2) ) {
	double mass    = FourMomentum( LeadingJet + SubLeadingJet ).mass()*GeV;
	double chi     = exp( fabs( LeadingJet.rapidity() - SubLeadingJet.rapidity() ) );
	double d_phi   = deltaPhi( LeadingJet.phi(), SubLeadingJet.phi() );
	double y_boost = 0.5 * (LeadingJet.rapidity() + SubLeadingJet.rapidity());

	_dijet_mass->fill(mass, weight);

	if (mass > 110) {
	  PhiCounter += weight;
	  _dijet_phi->fill(fabs(d_phi),weight);
	}
	if ( fabs(y_boost) < 1.1 ) {
	  if ( (mass > 110) && (mass < 370) ) {
	    ChiCounter1 += weight;
	    _dijet_chi_110_370->fill(chi,weight);
	  }
	  if ( (mass > 370) && (mass < 850) ) {
	    ChiCounter2 += weight;
	    _dijet_chi_370_850->fill(chi,weight);
	  }
	}
      }
    }

    void finalize() {
      
      //Scaling to cross section
      double xsec = crossSectionPerEvent()/(picobarn);
      //Normalizing to number of entries in histogram and mass 
      double chiScale1 = 1.0 / ( 260.0 * ChiCounter1 ); 
      double chiScale2 = 1.0 / ( 480.0 * ChiCounter2 ); 
      double phiScale  = 1.0 / ( (double) PhiCounter );
      
      //Additional Factors represent the division by rapidity
      _bjetpT_SV0.scale(xsec/2, this);

      scale(_bjetpT_SV0_All, xsec);
      scale(_bjetpT_pTRel,   xsec);
      scale(_dijet_mass,     xsec);
      scale(_dijet_phi        , phiScale );
      scale(_dijet_chi_110_370, chiScale1);
      scale(_dijet_chi_370_850, chiScale2);
    }

  private:
    
    BinnedHistogram<double> _bjetpT_SV0;

    AIDA::IHistogram1D * _bjetpT_SV0_All;
    AIDA::IHistogram1D * _bjetpT_pTRel;
    AIDA::IHistogram1D * _dijet_mass;
    AIDA::IHistogram1D * _dijet_phi;
    AIDA::IHistogram1D * _dijet_chi_110_370;
    AIDA::IHistogram1D * _dijet_chi_370_850;

    double ChiCounter1;
    double ChiCounter2;
    double PhiCounter;

  };

  // The hook for the plugin system
  DECLARE_RIVET_PLUGIN(ATLAS_2011_I930220);
}