// -*- C++ -*-
#include "Rivet/Analysis.hh"
#include "Rivet/Projections/FinalState.hh"

/// @todo Include more projections as required, e.g. ChargedFinalState, FastJets, ZFinder...
#include "Rivet/Projections/IdentifiedFinalState.hh"
#include "Rivet/Jet.hh"
#include "Rivet/Projections/FastJets.hh"

#include "fastjet/internal/base.hh"
#include "fastjet/JetDefinition.hh"
#include "fastjet/AreaDefinition.hh"
#include "fastjet/ClusterSequence.hh"
#include "fastjet/ClusterSequenceArea.hh"
#include "fastjet/PseudoJet.hh"

namespace Rivet {


  class ATLAS_2014_S1307756 : public Analysis {
  public:

    /// @name Constructors etc.
    //@{

    /// Constructor
    ATLAS_2014_S1307756()
      : Analysis("ATLAS_2014_S1307756")
    {      
      _eta_bins_areaoffset.push_back(0.0);
      _eta_bins_areaoffset.push_back(1.5);
      _eta_bins_areaoffset.push_back(3.0);
    }

    //@}


  public:

    /// @name Analysis methods
    //@{
    
    /// Book histograms and initialise projections before the run
    void init() {
      
      /// @todo Initialise and register projections here
      FinalState fs;
      addProjection(fs, "FS");

      FastJets fj(fs, FastJets::KT, 0.5);
      _area_def = new fastjet::AreaDefinition(fastjet::VoronoiAreaSpec());
      fj.useJetArea(_area_def);
      addProjection(fj, "KtJetsD05");

      IdentifiedFinalState photonfs(-2.37, 2.37, 22.0*GeV);
      photonfs.acceptId(PID::PHOTON);
      addProjection(photonfs, "photons");

      /// @todo initialize event count here:
      _Nevents = 0.;
      _fidCrossSection = 0.;
    }
    
    /// utility to compute ambiant energy density per eta bin
    int getEtaBin(double eta_w) const {
      double eta = fabs(eta_w);
      int v_iter=0;
      for (v_iter=0; v_iter < (int)_eta_bins_areaoffset.size()-1; v_iter++) {
        if (inRange(eta, _eta_bins_areaoffset[v_iter], _eta_bins_areaoffset[v_iter+1]))
          break;
      }
      return v_iter;
    }
    

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

      /// event count
      _Nevents += 1;

      /// event weight
      const double weight = event.weight();

      /// Require at least 2 photons in final state
      Particles photons = applyProjection<IdentifiedFinalState>(event, "photons").particlesByPt();
      if (photons.size() < 2) vetoEvent;
      
      /// compute the median energy density per eta bin
      std::vector<double> _ptDensity;
      std::vector< std::vector<double> > ptDensities;
      std::vector<double> emptyVec;
      ptDensities.assign(_eta_bins_areaoffset.size()-1, emptyVec);

      const fastjet::ClusterSequenceArea* clust_seq_area = applyProjection<FastJets>(event, "KtJetsD05").clusterSeqArea();
      foreach (const fastjet::PseudoJet& jet, applyProjection<FastJets>(event, "KtJetsD05").pseudoJets(0.0*GeV)) {
        double eta = fabs(jet.eta());
        double pt = fabs(jet.perp());

        /// get the cluster sequence
        double area = clust_seq_area->area(jet);

        if(area > 10e-4 && fabs(eta) < _eta_bins_areaoffset[_eta_bins_areaoffset.size()-1]) {
          ptDensities.at(getEtaBin(fabs(eta))).push_back(pt/area);
        }
      }

      for(int b=0; b<(int)_eta_bins_areaoffset.size()-1; b++) {
        double median = 0.0;
        if(ptDensities[b].size() > 0) {
            std::sort(ptDensities[b].begin(), ptDensities[b].end());
            int nDens = ptDensities[b].size();
            if( nDens%2 == 0 )
              median = (ptDensities[b][nDens/2] + ptDensities[b][(nDens-2)/2]) / 2;
            else
              median = ptDensities[b][(nDens-1)/2];
          }
        _ptDensity.push_back(median);
      }

      /// Loop over photons and find isolated ones
      Particles isolated_photons;

      foreach (const Particle& ph, photons) {
	Particles fs = applyProjection<FinalState>(event, "FS").particles();
	FourMomentum mom_in_EtCone;
	foreach (const Particle& p, fs) {
	  /// reject if the particle is not in DR=0.4 cone 
	  if (deltaR(ph.momentum(), p.momentum()) > 0.4) continue;
	  /// reject if the particle falls in the photon core 
	  if (fabs(ph.eta() - p.eta()) < 0.025 * 7 * 0.5 &&
	      fabs(ph.phi() - p.phi()) < PI/128. * 5 * 0.5) continue;
	  /// reject if the particle is a neutrino (muons are kept)
	  if(p.isNeutrino()) continue;
	  /// sum momentum
	  mom_in_EtCone += p.momentum();
	}

	/// subtract the UE correction (area*density)
        double EtCone_area = M_PI*.4*.4 - (7.0*.025)*(5.0*M_PI/128.);
        double correction = _ptDensity[getEtaBin(ph.eta())]*EtCone_area;

        /// Add isolated photon to list
        if (mom_in_EtCone.Et() - correction > 12*GeV) continue;
        isolated_photons.push_back(ph);
      }

      /// require at least two isolated photons
      if (isolated_photons.size() < 2)  vetoEvent ;    
   
      /// select leading pT pair
      std::sort(isolated_photons.begin(), isolated_photons.end(), cmpMomByPt);
      FourMomentum y1=isolated_photons[0].momentum();
      FourMomentum y2=isolated_photons[1].momentum();

      /// compute invariant mass 
      FourMomentum yy = y1+y2;
      double Myy = yy.mass();

      /// if Myy >= 110 GeV, apply relative cuts
      if (Myy/GeV >= 110 && (y1.Et()/Myy < 0.4 || y2.Et()/Myy < 0.3) ) vetoEvent;
                  
      /// Add to cross-section
      _fidCrossSection += weight;
    }


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

      // compute fiducial cross-section in fb
      double norm = crossSection()/sumOfWeights()/femtobarn;
      _fidCrossSection *= norm; 

      // compute selection efficiency & statistical error
      double eff = _fidCrossSection/(crossSection()/femtobarn) ;
      double err = sqrt(eff*(1-eff)/_Nevents);

      // Print out result
      Log& log = getLog(); 
      log << Log::INFO << "========================================" << endl;
      log << Log::INFO << "==== Total cross-section: " << crossSection()/femtobarn<< " fb" << endl;
      log << Log::INFO << "==== Fiducial cross-section: " << _fidCrossSection << " fb" << endl;
      log << Log::INFO << "========================================" << endl;
      log << Log::INFO << "==== Selection efficiency: " << eff << " +/- " << err << " (statistical error)" << endl;
      log << Log::INFO << "========================================" << endl;

    }

    //@}


  private:

    // Data members like post-cuts event weight counters go here

  private:

    fastjet::AreaDefinition* _area_def;
    std::vector<double> _eta_bins_areaoffset;
    float _fidCrossSection;
    int _Nevents;

  };



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

}