// -*- C++ -*-
#include "Rivet/Analysis.hh"
#include "Rivet/RivetAIDA.hh"
#include "Rivet/Tools/Logging.hh"
#include "Rivet/Projections/FinalState.hh"
#include "Rivet/Projections/ChargedFinalState.hh"
#include "Rivet/Projections/UnstableFinalState.hh"
#include "Rivet/Projections/MissingMomentum.hh"
#include "Rivet/Math/MathUtils.hh"
#include "LWH/Histogram1D.h"

namespace Rivet {


  namespace {

    inline double SumAB(vector<double>  vecX, vector<double> vecY) {
      double Sum(0) ;
      for (size_t i=0 ; i < vecX.size() ; i++) {
        Sum += vecX[i]*vecY[i] ;
      }
      return Sum ;
    }

    inline double SumA(vector<double>  vecX) {
      double Sum(0) ;
      for (size_t i=0 ; i < vecX.size() ; i++) {
        Sum += vecX[i] ;
      }
      return Sum ;
    }

    inline double Mean(vector<double>  vecX) {
      double mean(0) ;
      double N = (double) vecX.size() ;
      mean = SumA(vecX)/N ;
      return mean ;
    }

    inline double Standard_Deviation(vector<double> vecX) {
      double Sd(0) ;
      double N = (double) vecX.size() ;
      double Sum(0) ;
      double X_Bar = Mean(vecX) ;
      for (size_t i=0 ; i < vecX.size() ; i++) {
        Sum += (vecX[i]-X_Bar) * (vecX[i]-X_Bar) ;
      }
      Sd = sqrt(Sum/N) ;
      return Sd ;
    }

    inline double a0_regression(vector<double>  vecX, vector<double> vecY) {
      double a0(0),N(1000000),numerator(0),denumenator(10000) ;
      N = (double) vecX.size() ;
      numerator   = SumA(vecY) * SumAB(vecX,vecX) - SumA(vecX) * SumAB(vecX,vecY) ;
      denumenator = N * SumAB(vecX,vecX) - SumA(vecX) * SumA(vecX) ;
      a0 = numerator / denumenator ;
      return a0 ;
    }

    inline double a1_regression(vector<double>  vecX, vector<double> vecY) {
      double a1(0),N(10000000),numerator(0),denumenator(100000) ;
      N = (double) vecX.size() ;
      numerator   = N * SumAB(vecX,vecY) - SumA(vecX) * SumA(vecY) ;
      denumenator = N * SumAB(vecX,vecX) - SumA(vecX) * SumA(vecX) ;
      a1=numerator/ denumenator ;
      return a1 ;
    }

    inline double a1_regression2(vector<double>  vecX, vector<double> vecY) {
      double X_bar = Mean(vecX) ;
      double Y_bar = Mean(vecY) ;
      double SumXY(0) ;
      double N = (double) vecX.size() ;
      for (size_t i=0 ; i < vecX.size() ; i++) {
        SumXY += (vecY[i]-Y_bar) * (vecX[i]-X_bar) ;
      }
      double a1(0) ;
      a1 = SumXY / ( (Standard_Deviation(vecX) * Standard_Deviation(vecX)) * N) ;
      return a1 ;
    }

    /// @todo Unused?
    double Coefficient_of_correlation(vector<double>  vecX,vector<double>  vecY) {
      double X_bar = Mean(vecX) ;
      double Y_bar = Mean(vecY) ;
      double SumXY(0) ;
      double N = (double) vecX.size() ;
      for (size_t i=0 ; i < vecX.size() ; i++) {
        SumXY += (vecY[i]-Y_bar) * (vecX[i]-X_bar) ;
      }
      double rho(0) ; 
      rho= SumXY / ( (Standard_Deviation(vecX)*Standard_Deviation(vecY)) * N) ;
      return rho ; 
    }

    inline double Y_est(double X, double a0 , double a1) {
      return (a0+a1*X) ;
    }

    inline double Quadra_Sum_Residual(vector<double>  vecX,vector<double>  vecY) {
      double Sum(0) ;
      double Y_EST(0) ;
      double a0 = a0_regression(vecX,vecY) ;
      double a1 = a1_regression(vecX,vecY) ;
      for (size_t i=0 ; i < vecX.size() ; i++) {
        Y_EST = Y_est(vecX[i], a0 , a1)  ;
        Sum  += (vecY[i]-Y_EST) * (vecY[i]-Y_EST) ;
      }
      return Sum ;
    }

    inline double Error_On_Slope(vector<double>  vecX,vector<double>  vecY) {
      double eos(0) ;
      double N = (double) vecX.size() ;
      double quadra_sum_res = Quadra_Sum_Residual(vecX,vecY) ;
      double quadra_sum_x   = Standard_Deviation(vecX) * Standard_Deviation(vecX) * N ;

      eos = sqrt(quadra_sum_res/(N-2)) / sqrt(quadra_sum_x) ;
      return eos ;
    }

  }



  /// Generic analysis looking at various distributions of final state particles
  class ATLAS_2012_I1093734 : public Analysis {
  public:

    /// Constructor
    ATLAS_2012_I1093734()
      : Analysis("ATLAS_2012_I1093734")
    {
      // Stat convergence happens around 20k events, so it doesn't make sense to run this
      // analysis with much less than that. Given that, lets avoid some unnecessary vector
      // resizing by allocating sensible amounts in the first place.
      _vecWeight.reserve(10000);
      
      for (int ipt = 0; ipt < NptBins; ++ipt) {
	
	for (int k = 0; k < NetaBins; ++k) {
	  _vecsNchF[ipt][k].reserve(10000);
	  _vecsNchB[ipt][k].reserve(10000);
	  
	  if (ipt == 0) {
	    _vecsSumptF[k].reserve(10000);
	    _vecsSumptB[k].reserve(10000);
	  }
	}
      }
    }


  public:

    /// @name Analysis methods
    //@{

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

      //----------------------
      //FB correlations part
      //----------------------

      // Projections
      for (int ipt = 0; ipt < NptBins; ++ipt) {
	
        double ptmin = ptminValue[ipt] * MeV;
        
        for (int ieta = 0; ieta < NetaBins; ++ieta) {
	  addProjection(ChargedFinalState(-etaValue[ieta]    , -etaValue[ieta]+0.5, ptmin), "Tracks"+etabinnames[ieta]+"B"+ptbinnames[ipt]);
   	  addProjection(ChargedFinalState( etaValue[ieta]-0.5,  etaValue[ieta]    , ptmin), "Tracks"+etabinnames[ieta]+"F"+ptbinnames[ipt]);
	}
	
        const string chname = "CFS" + ptbinnames[ipt];
	addProjection( ChargedFinalState(-2.5,2.5,ptmin), chname );
      }
      
      if      (fuzzyEquals(sqrtS()*GeV, 7000, 1E-3)) {
	for (int ipt  = 0; ipt  < NptBins ; ++ipt )  _histNchCorr_vsEta[ipt ]= bookDataPointSet( 1+ipt ,2,1);
	for (int ieta = 0; ieta < NetaBins; ++ieta)  _histNchCorr_vsPt [ieta]= bookDataPointSet( 8+ieta,2,1);
	_histPtsumCorr = bookDataPointSet(13,2,1);
      }
      else if (fuzzyEquals(sqrtS()*GeV,  900, 1E-3)) {
	_histNchCorr_vsEta[0] = bookDataPointSet(14,2,1);
	_histPtsumCorr        = bookDataPointSet(15,2,1);
      }
      
      
      //----------------------
      //Azimutal correlations part
      //----------------------
      double ptmin = 500*MeV;
      addProjection(ChargedFinalState(-2.5, 2.5, ptmin), "ChargedTracks25");
      addProjection(ChargedFinalState(-2.0, 2.0, ptmin), "ChargedTracks20");
      addProjection(ChargedFinalState(-1.0, 1.0, ptmin), "ChargedTracks10");

      h_dphi[0] = bookHistogram1D("dphi0", 50, 0.0, PI);
      h_same[0] = bookHistogram1D("same0", 50, 0.0, PI);
      h_oppo[0] = bookHistogram1D("oppo0", 50, 0.0, PI);
      h_dphi[1] = bookHistogram1D("dphi1", 50, 0.0, PI);
      h_same[1] = bookHistogram1D("same1", 50, 0.0, PI);
      h_oppo[1] = bookHistogram1D("oppo1", 50, 0.0, PI);
      h_dphi[2] = bookHistogram1D("dphi2", 50, 0.0, PI);
      h_same[2] = bookHistogram1D("same2", 50, 0.0, PI);
      h_oppo[2] = bookHistogram1D("oppo2", 50, 0.0, PI);

      if      (fuzzyEquals(sqrtS()*GeV, 7000, 1E-3)) {
        h_dphiMin[0]       = bookDataPointSet(2,1,1);
        h_dphiMin[1]       = bookDataPointSet(4,1,1);
        h_dphiMin[2]       = bookDataPointSet(6,1,1);
      }
      else if (fuzzyEquals(sqrtS()*GeV,  900, 1E-3)) {
        h_dphiMin[0]       = bookDataPointSet(1,1,1);
        h_dphiMin[1]       = bookDataPointSet(3,1,1);
        h_dphiMin[2]       = bookDataPointSet(5,1,1);
      }
    }




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

      double weight = event.weight();

      bool runAzimCorrPart=true;
      
      
      for (int ipt = 0; ipt < NptBins; ++ipt) {

	const string chname = "CFS" + ptbinnames[ipt];

	const ChargedFinalState& charged = applyProjection<ChargedFinalState>(event, chname);
	
	if (charged.particles().size()>=2) {
	  
	  if (ipt == 0) _vecWeight.push_back(weight);
	  
	  for (int ieta = 0; ieta < NetaBins; ++ieta) {
	    
	    const string fname = "Tracks" + etabinnames[ieta] + "F" + ptbinnames[ipt];
	    const string bname = "Tracks" + etabinnames[ieta] + "B" + ptbinnames[ipt];
	    
	    const ParticleVector particlesF = applyProjection<FinalState>(event, fname).particles();
	    const ParticleVector particlesB = applyProjection<FinalState>(event, bname).particles();
	    
	    _vecsNchF  [ipt][ieta].push_back((double)particlesF.size());
	    _vecsNchB  [ipt][ieta].push_back((double)particlesB.size());
	    
	    //sum pt only for 100 MeV particles
	    if (ipt == 0) {

	      double sumptF = 0;
	      double sumptB = 0;

	      foreach (const Particle& p, particlesF) sumptF += p.momentum().pT();
	      foreach (const Particle& p, particlesB) sumptB += p.momentum().pT();
	      
	      _vecsSumptF[ieta].push_back(sumptF);
	      _vecsSumptB[ieta].push_back(sumptB);
	    }
	  }
	}
      }
      

      if (runAzimCorrPart) {
	
	string etabin[3] = { "10", "20", "25" };
	
	for (int iColl=0; iColl<3; iColl++) {

	  const string fname = "ChargedTracks" + etabin[iColl] ;

	  const ParticleVector partTrk = applyProjection<FinalState>(event, fname).particles();

	  //First loop over particles to find the leading track
	  double leadpt  =  0;
	  double leadphi = -99;
	  double leadeta = -99;
	  int leadindex  = -1;
	  int count      =  0;
	  
	  for (std::vector<Particle>::const_iterator it = partTrk.begin(); it != partTrk.end(); ++it)  {

	    const Particle p=*it;

	    if ( p.momentum().pT() > leadpt)  {
	      leadphi   = p.momentum().phi();
	      leadeta   = p.momentum().eta();
	      leadpt    = p.momentum().pT();
	      leadindex = it - partTrk.begin();
	    }
	    count++;
	  }
	  
	  //Second loop to fill the histograms
	  for (std::vector<Particle>::const_iterator it = partTrk.begin(); it != partTrk.end(); ++it)  {

	    if ( (it-partTrk.begin()) != leadindex) {
	      const Particle p= *it;
	      double dphi = deltaPhi( p.momentum().phi(), leadphi);
	      //Fill histograms
	      h_dphi[iColl]->fill(dphi, weight);
	      float sign = leadeta * p.momentum().eta();
	      if (sign>0) h_same[iColl]->fill(dphi, weight);
	      if (sign<0) h_oppo[iColl]->fill(dphi, weight);
	    }
	  }
	}
      }
    }


    /// Finalize
    void finalize() {

      //-----------------
      //FB part
      //-----------------
      // @todo for 2D plots will need _vecsNchF[i], _vecsNchB[j]
      for (int ipt=0; ipt < NptBins; ++ipt) {

	//just one plot at 900 GeV
	if ( fuzzyEquals(sqrtS()*GeV,  900, 1E-3) && (ipt > 0) ) continue;

	vector<double> nchval, ncherr;
	
	for (int k = 0; k < NetaBins; ++k) {
	  nchval.push_back( a1_regression2( _vecsNchF[ipt][k], _vecsNchB[ipt][k]) );
	  ncherr.push_back( Error_On_Slope( _vecsNchF[ipt][k], _vecsNchB[ipt][k]) );
	}
	
	_histNchCorr_vsEta[ipt]->setCoordinate(1, nchval, ncherr);
      }
      
      for (int ieta = 0; ieta < NetaBins; ++ieta) {

	//no plots at 900 GeV
	if ( fuzzyEquals(sqrtS()*GeV,  900, 1E-3) ) continue;
	
	vector<double> nchval, ncherr;
	
	for (int ipt = 0; ipt < NptBins; ++ipt) {
	  nchval.push_back( a1_regression2( _vecsNchF[ipt][ieta], _vecsNchB[ipt][ieta]) );
	  ncherr.push_back( Error_On_Slope( _vecsNchF[ipt][ieta], _vecsNchB[ipt][ieta]) );
	}
	
	_histNchCorr_vsPt [ieta]->setCoordinate(1, nchval, ncherr);
      }
      
      vector<double> ptsval, ptserr;
      for (int k = 0; k < NetaBins; ++k) {
        ptsval.push_back(a1_regression2(_vecsSumptF[k], _vecsSumptB[k]));
        ptserr.push_back(Error_On_Slope(_vecsSumptF[k], _vecsSumptB[k]));
      }
      _histPtsumCorr->setCoordinate(1, ptsval, ptserr);

      //-----------------
      //azimuthal part
      //-----------------
      AIDA::IHistogramFactory& hf = histogramFactory();

      AIDA::IHistogram1D* h_SameMinusOppo[3];
      
      for (int iColl=0; iColl<3; iColl++) {
	
        if      (fuzzyEquals(sqrtS()*GeV, 7000, 1E-3)) {
          h_SameMinusOppo[iColl] = hf.subtract( histoPath(8+2*iColl,1,1), *h_same[iColl], *h_oppo[iColl] );
          scale(h_SameMinusOppo[iColl], 1./h_SameMinusOppo[iColl]->sumBinHeights() * h_SameMinusOppo[iColl]->axis().binWidth(1) );
        }
        else if (fuzzyEquals(sqrtS()*GeV,  900, 1E-3)) {
          h_SameMinusOppo[iColl] = hf.subtract( histoPath(7+2*iColl,1,1), *h_same[iColl], *h_oppo[iColl] );
          scale(h_SameMinusOppo[iColl], 1./h_SameMinusOppo[iColl]->sumBinHeights() * h_SameMinusOppo[iColl]->axis().binWidth(1) );
        }
	
	int    Nbins   = h_dphi[iColl]->axis().bins();
	double histMin = h_dphi[iColl]->minBinHeight();
	
	std::vector<double> diff (Nbins, 0.0);
	std::vector<double> err  (Nbins, 0.0);
	
        //subtract minimal value
        double sumDiff = 0.;
	for (int n = 0; n < Nbins; ++n) {
	  diff[n] = h_dphi[iColl]->binHeight(n) - histMin;
	  err[n]  = h_dphi[iColl]->binError(n);
	  sumDiff += diff[n];
	}
        //normalize
	for (int n = 0; n < Nbins; ++n) {
	  diff[n] = diff[n]/sumDiff;
	  err[n]  = err[n]/sumDiff;
	}
	
	h_dphiMin[iColl]->setCoordinate(1,diff,err);
      } 
      
      
      // remove temporary histograms
      for (size_t i=0; i<3; i++) {
        hf.destroy(h_dphi[i]);
        hf.destroy(h_same[i]);
        hf.destroy(h_oppo[i]);
      }
      
    }

    //@}

    
  private:

    //----------------------
    //FB correlations part
    //----------------------
    static const int NptBins  = 7;
    static const int NetaBins = 5;

    static const double ptminValue[NptBins];
    static const string ptbinnames[NptBins];

    static const double etaValue   [NetaBins];
    static const string etabinnames[NetaBins];


    vector<double> _vecWeight;
    vector<double> _vecsNchF[NptBins][NetaBins];
    vector<double> _vecsNchB[NptBins][NetaBins];
    vector<double> _vecsSumptF[NetaBins];
    vector<double> _vecsSumptB[NetaBins];
    
    /// @name Histograms
    AIDA::IDataPointSet *_histNchCorr_vsEta[NptBins ]; 
    AIDA::IDataPointSet *_histNchCorr_vsPt [NetaBins]; 
    AIDA::IDataPointSet *_histPtsumCorr;
    
    //----------------------
    //Azimutal correlations part
    //----------------------
    AIDA::IHistogram1D* h_dphi[3];
    AIDA::IHistogram1D* h_same[3];
    AIDA::IHistogram1D* h_oppo[3];
    AIDA::IDataPointSet* h_dphiMin[3];
  };

  const double ATLAS_2012_I1093734::ptminValue[] = {100, 200, 300, 500, 1000, 1500, 2000 };
  const string ATLAS_2012_I1093734::ptbinnames[] = { "100", "200", "300", "500", "1000", "1500", "2000" };
  
  const double ATLAS_2012_I1093734::etaValue   [] = {0.5, 1.0, 1.5, 2.0, 2.5};
  const string ATLAS_2012_I1093734::etabinnames[] = { "05", "10", "15", "20", "25" };


  // This global object acts as a hook for the plugin system
  DECLARE_RIVET_PLUGIN(ATLAS_2012_I1093734);

}