// -*- C++ -*-
#include "Rivet/Analysis.hh"
#include "Rivet/RivetAIDA.hh"
#include "Rivet/Tools/Logging.hh"
#include "Rivet/Projections/FastJets.hh"
#include "Rivet/Projections/FinalState.hh"
/// @todo Include more projections as required, e.g. ChargedFinalState, FastJets, ZFinder...

namespace Rivet {


  class ATLAS_2012_I1119557 : public Analysis {
  public:

    ATLAS_2012_I1119557()
      : Analysis("ATLAS_2012_I1119557")
    {    }


  public:

    double jetWidth(const Jet& jet) {
      double phi_jet = jet.phi();
      double eta_jet = jet.eta();
      double width = 0.0;
      double pTsum = 0.0;
      foreach (const Particle& p, jet.particles()) {
        double pT = p.momentum().pT();
        double eta = p.momentum().eta();
        double phi = p.momentum().phi();
        width += sqrt(pow(phi_jet - phi,2) + pow(eta_jet - eta ,2)) * pT;
        pTsum += pT;
      }
      if (pTsum != 0.0)  return width/pTsum;
      else               return -1;
    }

    // This is the code for the eccentricity calculation, copied and adapted from Lily's code
    double getEcc(const Jet& jet) {

      vector<double> phis;
      vector<double> etas;
      vector<double> energies;

      double etaSum = 0.; double phiSum = 0.; double eTot = 0.; 
      foreach (const Particle& p, jet.particles()) {

        double E = p.momentum().E();
        double eta = p.momentum().eta();
        
        energies.push_back(E);
        etas.push_back(jet.momentum().eta() - eta);
        
        eTot   += E;
        etaSum += eta * E;

        double dPhi = mapAngleMPiToPi( jet.momentum().phi() - p.momentum().phi() );

        //if DPhi does not lie within 0 < DPhi < PI take 2*PI off DPhi
        //this covers cases where DPhi is greater than PI
        //if( fabs( dPhi - TWOPI ) < fabs(dPhi) ) dPhi -= TWOPI;
        //if DPhi does not lie within -PI < DPhi < 0 add 2*PI to DPhi
        //this covers cases where DPhi is less than -PI
        //else if( fabs(dPhi + TWOPI) < fabs(dPhi) ) dPhi += TWOPI;
       
        phis.push_back(dPhi);
        
        phiSum += dPhi * E;
      }

      //these are the "pull" away from the jet axis
      etaSum = etaSum/eTot; 
      phiSum = phiSum/eTot;

      // now for every cluster we alter its position by moving it:
      // away from the new axis if it is in the direction of -ve pull
      // closer to the new axis if it is in the direction of +ve pull
      // the effect of this will be that the new energy weighted center will be on the old jet axis. 
      double little_x=0., little_y=0.;
      for(unsigned int k = 0; k< jet.particles().size(); k++) {
        little_x+= etas[k]-etaSum;
        little_y+= phis[k]-phiSum;

        etas[k] = etas[k]-etaSum;
        phis[k] = phis[k]-phiSum;
     }

      double X1=0.; 
      double X2=0.; 
      for(unsigned int i = 0; i < jet.particles().size(); i++) {
        X1 += 2. * energies[i]* etas[i] * phis[i]; // this is =2*X*Y
        X2 += energies[i]*(phis[i] * phis[i] - etas[i] * etas[i] ); // this isX^2 - Y^2
      }
      
        // variance calculations 
      double Theta = .5*atan2(X1,X2);

      double sinTheta =sin(Theta);
      double cosTheta = cos(Theta);
      double Theta2 = Theta + 0.5*PI;
      double sinThetaPrime = sin(Theta2);
      double cosThetaPrime = cos(Theta2);

      double VarX = 0.;
      double VarY = 0.;
      for(unsigned int i = 0; i < jet.particles().size(); i++) {
        double X=sinTheta*etas[i] + cosTheta*phis[i];    
        double Y=sinThetaPrime*etas[i] + cosThetaPrime*phis[i]; 
        VarX += energies[i]*X*X;
        VarY += energies[i]*Y*Y;
      }
       
      double VarianceMax = VarX; 
      double VarianceMin = VarY;
     
      if(VarianceMax < VarianceMin) {
        VarianceMax = VarY;
        VarianceMin = VarX;
      }
      
      double ECC;
      if (VarianceMax != 0.0)  ECC = 1.0 - (VarianceMin/VarianceMax);
      else                     ECC = -1.0;

      return ECC;
    }

    // This is the code for the planar flow calculation, copied and adapted from Lily's code
    double getPFlow(const  Jet& jet) {

      double phi0=jet.momentum().phi();
      double eta0=jet.momentum().eta();

      double nref[3];
      nref[0]=(cos(phi0)/cosh(eta0));
      nref[1]=(sin(phi0)/cosh(eta0));
      nref[2]=tanh(eta0);
      
      // This is the rotation matrix
      double RotationMatrix[3][3];
      CalcRotationMatrix(nref, RotationMatrix);
      
      double Iw00(0.), Iw01(0.), Iw11(0.), Iw10(0.);
     
      foreach (const Particle& p, jet.particles()) {
        double a=1./(p.momentum().E()*jet.momentum().mass());
        FourMomentum rotclus = RotateAxes(p.momentum(), RotationMatrix);
        Iw00 += a*pow(rotclus.px(), 2);
        Iw01 += a*rotclus.px()*rotclus.py();
        Iw10 += a*rotclus.py()*rotclus.px();
        Iw11 += a*pow(rotclus.py(), 2);
      }

      double det=Iw00*Iw11-Iw01*Iw10;
      double trace=Iw00+Iw11;

      double pf;
      if (trace != 0.0) pf=(4.0*det)/(pow(trace,2));
      else              pf=-1.0;

      return pf;
    }
   

    // This is the code for the angularity calculation, copied and adapted from Lily's code
    double getAngularity(const Jet& jet) {
      double sum_a=0.;
      //This a used in angularity calc can take any value <2 (e.g. 1,0,-0.5 etc) for infrared safety
      const double a=-2.;
     
      foreach (const Particle& p, jet.particles()) {
        double e_i       = p.momentum().E();
        double theta_i   = jet.momentum().angle(p.momentum());
        double e_theta_i = e_i * pow(sin(theta_i),a) * pow(1-cos(theta_i),1-a);
        sum_a += e_theta_i;
      }

      double Angularity;
      if (jet.momentum().mass() != 0.0)  Angularity = sum_a/jet.momentum().mass();//mass is in MeV
      else                             Angularity = -1.0;

      return Angularity;
    }

    // Adapted code from Lily  
    FourMomentum RotateAxes(const Rivet::FourMomentum& p, double M[3][3]){
      double px_rot=M[0][0]*(p.px())+M[0][1]*(p.py())+M[0][2]*(p.pz());
      double py_rot=M[1][0]*(p.px())+M[1][1]*(p.py())+M[1][2]*(p.pz());
      double pz_rot=M[2][0]*(p.px())+M[2][1]*(p.py())+M[2][2]*(p.pz()); 
      return FourMomentum(p.E(), px_rot, py_rot, pz_rot);
    } 
    
    // Untouched code from Lily  
    void CalcRotationMatrix(double nvec[3],double rot_mat[3][3]){
      //clear momentum tensor
      for(int i=0; i<3; i++){
        for(int j=0; j<3; j++){
          rot_mat[i][j]=0.;
        }
      }
      double mag3=sqrt(nvec[0]*nvec[0] + nvec[1]*nvec[1]+ nvec[2]*nvec[2]);
      double mag2=sqrt(nvec[0]*nvec[0] + nvec[1]*nvec[1]);
      if(mag3<=0){
        cout<<"rotation axis is null"<<endl;
      }
      assert(mag3 > 0);

      double ctheta0=nvec[2]/mag3;
      double stheta0=mag2/mag3;
      double cphi0 = (mag2>0.) ? nvec[0]/mag2:0.;
      double sphi0 = (mag2>0.) ? nvec[1]/mag2:0.;

      rot_mat[0][0]=-ctheta0*cphi0;
      rot_mat[0][1]=-ctheta0*sphi0;
      rot_mat[0][2]=stheta0;
      rot_mat[1][0]=sphi0;
      rot_mat[1][1]=-1.*cphi0;
      rot_mat[1][2]=0.;
      rot_mat[2][0]=stheta0*cphi0;
      rot_mat[2][1]=stheta0*sphi0;
      rot_mat[2][2]=ctheta0;

      return;
    }
    
    void init() {

      const FinalState fs;
      addProjection(fs,"FinalState");
      
      FastJets fj06(fs,  FastJets::ANTIKT, 0.6);
      //fj06.useInvisibles();
      addProjection(fj06, "AntiKT06");

      FastJets fj10(fs,  FastJets::ANTIKT, 1.0);
      //fj10.useInvisibles();
      addProjection(fj10, "AntiKT10");


      for (size_t alg = 0; alg < 2; ++alg) {
        _Mass[alg]         = bookHistogram1D(1, alg+1, 1);
        _Width[alg]        = bookHistogram1D(2, alg+1, 1);
	//commented out for now
        //_Eccentricity[alg] = bookHistogram1D(3, alg+1, 1);
      }
      _PlanarFlow = bookHistogram1D(4, 2, 1);
      _Angularity = bookHistogram1D(5, 1, 1);

    }


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

      const double weight = event.weight();
      Jets jetAr[2];
      jetAr[0] = applyProjection<FastJets>(event, "AntiKT06").jetsByPt(300.*GeV, MAXDOUBLE, -2.0, 2.0);
      jetAr[1] = applyProjection<FastJets>(event, "AntiKT10").jetsByPt(300.*GeV, MAXDOUBLE, -2.0, 2.0);
      
      for (size_t alg = 0; alg < 2; ++alg) {

        // Require at least one jet
        if (jetAr[alg].size() < 1) continue;

        // The leading jet
        const Jet& jet = jetAr[alg][0];
        double m   = jet.momentum().mass()/GeV;
        double eta = jet.momentum().eta();

        _Mass[alg]   ->fill(m, weight);
        _Width[alg]  ->fill(jetWidth(jet), weight);

        // Eccentricity -commented out for now
        //if (fabs(eta) < 0.7 && m > 100.) {
        //  _Eccentricity[alg]->fill(getEcc(jet), weight);
        //}

        if (fabs(eta) < 0.7 && m > 130. && m < 210. && alg==1) _PlanarFlow->fill(getPFlow(jet), weight);
        if (fabs(eta) < 0.7 && m > 100. && m < 130. && alg==0) _Angularity->fill(getAngularity(jet), weight);
      } // Close algorithm loop
    }


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

      for (size_t alg = 0; alg < 2; ++alg) {
        normalize(_Mass[alg]);
        normalize(_Width[alg]);
	//normalize(_Eccentricity[alg]);
      }
      normalize(_PlanarFlow);
      normalize(_Angularity);
 
    }

  private:

    AIDA::IHistogram1D *_Mass[2];
    AIDA::IHistogram1D *_Width[2];
    //AIDA::IHistogram1D *_Eccentricity[2];
    AIDA::IHistogram1D *_PlanarFlow;
    AIDA::IHistogram1D *_Angularity;
  };



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

}