// JetsWithoutJets Package
// Questions/Comments? danbert@mit.edu jthaler@mit.edu
//
// Copyright (c) 2013
// Daniele Bertolini and Jesse Thaler
//
// $Id: $
//----------------------------------------------------------------------
// This file is part of FastJet contrib.
//
// It is free software; you can redistribute it and/or modify it under
// the terms of the GNU General Public License as published by the
// Free Software Foundation; either version 2 of the License, or (at
// your option) any later version.
//
// It is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
// or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
// License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this code. If not, see <http://www.gnu.org/licenses/>.
//----------------------------------------------------------------------

#include "JetsWithoutJets.hh"
#include <sstream>

using namespace std;


FASTJET_BEGIN_NAMESPACE      // defined in fastjet/internal/base.hh

namespace contrib {

//////
//
// Storage Array
//
//////


// Creates storage array
void JWJStorageArray::establishStorage(const vector<fastjet::PseudoJet> & my_jets, double Rjet, double ptcut) {
   
   // set radius and ptcut
   _Rjet = Rjet;
   _ptcut = ptcut;
   FunctorJetScalarPt sumPt = FunctorJetScalarPt();

   // Dividing up Phase Space to bins of (approximate) size 2R by 2R
   // Actually makes bins of size bigger than 2R by 2R such that /
   // we have equal spacing of regions.
   // (This index scheme need to be synced with getRapIndex/getPhiIndex)
   _maxRapIndex = (int) floor((_rapmax)/_Rjet);// 2 rapmax / 2 Rjet
   _rapSpread = 2.0*_rapmax/(floor((_rapmax)/_Rjet));
   _maxPhiIndex = (int) floor((M_PI)/_Rjet); // 2 pi / 2 Rjet
   _phiSpread = 2.0*M_PI/(floor((M_PI)/_Rjet));
     
   // Initializing storage
   _regionStorage.resize(_maxRapIndex);
   _aboveCutBool.resize(_maxRapIndex);
   for (int rapIndex = 0 ; rapIndex < _maxRapIndex ; rapIndex++) {
      _regionStorage[rapIndex].resize(_maxPhiIndex);
      _aboveCutBool[rapIndex].resize(_maxPhiIndex);
      for (int phiIndex = 0; phiIndex < _maxPhiIndex; phiIndex++) {
         _regionStorage[rapIndex][phiIndex].clear();
      }
   }

   // looping through particles and assigning to bins 
   for (unsigned int i = 0; i < my_jets.size(); i++) {
      PseudoJet currentJet = my_jets[i];
      double rap = currentJet.rap();
      double phi = currentJet.phi();

      // find index for particle
      int lowRap = (int) floor((rap + _rapmax)/_rapSpread);
      int highRap = (int) ceil((rap + _rapmax)/_rapSpread);
      int lowPhi = (int) floor(phi/_phiSpread);
      int highPhi = (int) ceil(phi/_phiSpread);
      if (highPhi >= _maxPhiIndex) highPhi = highPhi - _maxPhiIndex; // loop around

      // if out of bounds, set to be in bounds.
      if (lowRap < 0) lowRap = 0;
      if (lowRap >= _maxRapIndex) lowRap = _maxRapIndex-1;
      if (highRap < 0) highRap = 0;
      if (highRap >= _maxRapIndex) highRap = _maxRapIndex-1;

      // Fill in storage.  For particles at periphery in rapidity, only fill two bins instead of four.  If phi overlaps, do the same thing.
      _regionStorage[lowRap][lowPhi].push_back(currentJet);
      if (lowPhi != highPhi) _regionStorage[lowRap][highPhi].push_back(currentJet);
      if (lowRap != highRap) _regionStorage[highRap][lowPhi].push_back(currentJet);
      if (lowRap != highRap && lowPhi != highPhi) _regionStorage[highRap][highPhi].push_back(currentJet);

   }

   // store whether above ptCut values
   for (int rapIndex = 0; rapIndex < _maxRapIndex; rapIndex++) {
      for (int phiIndex = 0; phiIndex < _maxPhiIndex; phiIndex++) {
         _aboveCutBool[rapIndex][phiIndex] = (sumPt(_regionStorage[rapIndex][phiIndex]) >= ptcut);
      }
   }

}

// What is the rapidity index?  (Sync with JWJStorageArray::establishStorage)
int JWJStorageArray::getRapIndex(const fastjet::PseudoJet & currentJet) const {
   double rap = currentJet.rap();
   int rapIndex = round((rap + _rapmax)/_rapSpread);

   if (rapIndex < 0) rapIndex = 0;
   if (rapIndex >= _maxRapIndex) rapIndex = _maxRapIndex - 1;

   return rapIndex;
}

// what is the phi index? (Sync with JWJStorageArray::establishStorage)
int JWJStorageArray::getPhiIndex(const fastjet::PseudoJet & currentJet) const {
   double phi = currentJet.phi();
   int phiIndex = round(phi/_phiSpread);
   
   if (phiIndex >= _maxPhiIndex) phiIndex = phiIndex - _maxPhiIndex;
   
   return phiIndex;
}

vector<fastjet::PseudoJet> & JWJStorageArray::getStorageFor(const fastjet::PseudoJet & currentJet) {
   return _regionStorage[getRapIndex(currentJet)][getPhiIndex(currentJet)];
}

bool JWJStorageArray::aboveCutFor(const fastjet::PseudoJet & currentJet) {
   return _aboveCutBool[getRapIndex(currentJet)][getPhiIndex(currentJet)];
}


//////
//
// Extendable Jet-like Event Shape
//
//////


// If called, this helper function establishes the storage array for speed up
void JetLikeEventShape::establishStorage(const std::vector<PseudoJet> & particles) const {
   if (_useStorageArray) {
      _myStorageArray.establishStorage(particles,_Rjet, _ptcut);
   } else {
      // do nothing
   }
}

// This helper function calculates the weights assigned to each particle.
// In most cases, the user does not need to access the information set by this function
void JetLikeEventShape::establishWeights(const PseudoJet myParticle, const std::vector<PseudoJet> & particles) const {
      
      FunctorJetScalarPt sumPt = FunctorJetScalarPt();
      _includeParticle = false;
      _weightParticle = 0.0;
      _pt_in_Rjet = 0.0;
      _pt_in_Rsub = 0.0;
            
      // Find particles in my _Rjet neighborhood
      fastjet::Selector sel = fastjet::SelectorCircle(_Rjet);
      sel.set_reference(myParticle);

      if (_useStorageArray) { // start from rapidity/phi blocks
         if (!_myStorageArray.aboveCutFor(myParticle)) {
            return; //don't do any further analysis if not above jet pTcut
         } else {
            _nearbyParticles = sel(_myStorageArray.getStorageFor(myParticle));
         } 
      } else { // use all particles
         _nearbyParticles = sel(particles);
      }
      // See if there is enough pt in the neighborhood to pass _ptcut
      _pt_in_Rjet = sumPt(_nearbyParticles);
      if (_pt_in_Rjet < _ptcut) return;
      
      // If trimming is on, do check of _fcut as well
      if (_trim) {
         assert(_Rsub <= _Rjet);
         fastjet::Selector sel_sub = fastjet::SelectorCircle(_Rsub);
         sel_sub.set_reference(myParticle);
         std::vector<fastjet::PseudoJet> near_particles_sub = sel_sub(_nearbyParticles); // Save time and only look at near particles 
         _pt_in_Rsub = sumPt(near_particles_sub);
         if (_pt_in_Rsub/_pt_in_Rjet < _fcut) return;
      }
                  
      // If enough pt, find weight and apply measurement
      _includeParticle = true;
      _weightParticle = myParticle.pt()/_pt_in_Rjet;
      
}

// For a generic event shape, the result is obtained by simply multiplying the weight
// of a given particle by the _measurement acting on the _nearbyParticles
double JetLikeEventShape::result(const std::vector<PseudoJet> & particles) const {
   
   double answer = 0.0;
   establishStorage(particles);
  
   for (unsigned int i = 0; i < particles.size(); i++) {            
      establishWeights(particles[i],particles);
      if (_includeParticle) {
         double measurement = _measurement->result(_nearbyParticles);
         answer += measurement*_weightParticle;
      }
   }

   return(answer);
}



//////
//
// MissingTransverseMomentum
//
//////

// To calculate missing pT, we cannot use a standard _measurement, so we code this by hand.
double ShapeMissingPt::result(const std::vector<PseudoJet> & particles) const {
   
   PseudoJet p_tot(0,0,0,0);
   establishStorage(particles);

  
   for (unsigned int i = 0; i < particles.size(); i++) {            
      establishWeights(particles[i],particles);
      if (_includeParticle) {
         p_tot += particles[i];
      }
   }

   return(p_tot.pt());
}

string ShapeMissingPt::description() const {
  return "Missing Transverse Momentum as event shape, " + jetParameterString();
}


//////
//
// TrimmedSubjetMultiplicity
//
//////

//  While it would be possible to calculate trimmed subjet multiplicity using a
//  _measurement functor, it is computationally more efficient to use the information
// already captured by establishWeights().
double ShapeTrimmedSubjetMultiplicity::result(const vector<PseudoJet> & particles) const {

   double answer = 0.0;
   establishStorage(particles);

   for (unsigned int i = 0; i < particles.size(); i++) {            
      establishWeights(particles[i],particles);
      if (_includeParticle && _pt_in_Rsub > _ptsubcut) {
         answer += particles[i].pt()/_pt_in_Rsub;
      }
   }
   return(answer);

}

string ShapeTrimmedSubjetMultiplicity::description() const {
  ostringstream oss;
  oss << "Trimmed Subjet multiplicity with R_jet=" << _Rjet << ", pT_cut=" << _ptcut << ", R_sub=" << _Rsub << ", and fcut=" << _fcut; 
  return oss.str();
}

//////
//
// Shape Trimming
//
//////

//----------------------------------------------------------------------
// Helper for SelectorShapeTrimming. 
// Class derived from SelectorWorker: pass, terminator, applies_jet_by_jet, and description are overloaded. 
// It can't be applied to an individual jet since it requires information about the neighbourhood of the jet.
class SW_ShapeTrimming: public SelectorWorker {

private:

   double _Rjet, _ptcut, _Rsub, _fcut;

   bool _useTempStorage;

public:
   
   SW_ShapeTrimming(double Rjet, double ptcut, double Rsub, double fcut, bool useTempStorage = true) : _Rjet(Rjet), _ptcut(ptcut), _Rsub(Rsub), _fcut(fcut), _useTempStorage(useTempStorage) {};

   virtual bool pass(const PseudoJet &) const {
	if (!applies_jet_by_jet())	
     	throw Error("Cannot apply this selector worker to an individual jet");
      return false;
   }

   virtual void terminator(vector<const PseudoJet *> & jets) const {
      
      // Copy pointers content to a vector of PseudoJets. Check if the pointer has been already nullified. 
      vector<PseudoJet> my_jets;
      vector<unsigned int> indices; 
      for (unsigned int i=0; i < jets.size(); i++){
         if (jets[i]){
            indices.push_back(i);
            my_jets.push_back(*jets[i]);
         } 	
      }
	   
      FunctorJetScalarPt sumPt = FunctorJetScalarPt();
      JWJStorageArray myStorageArray = JWJStorageArray();
      
      // to increase speed in very high multiplicity events, make a grid of rapidity/phi blocks
      if (_useTempStorage) {
         myStorageArray.establishStorage(my_jets,_Rjet,_ptcut);
      }
   
   	for (unsigned int i=0; i < my_jets.size(); i++){
         
         // Select particles in radius _Rjet
         fastjet::Selector sel = fastjet::SelectorCircle(_Rjet);
         sel.set_reference(my_jets[i]);
         vector<PseudoJet> near_particles;
     
         if (_useTempStorage) { // start from rapidity/phi blocks
            if (!myStorageArray.aboveCutFor(my_jets[i])) {
               jets[indices[i]] = NULL;
               continue; // no need to consider that particle further
            } else {
               near_particles = sel(myStorageArray.getStorageFor(my_jets[i]));
            }
         } else { // use all particles
            near_particles = sel(my_jets);
         }

         // Calculate enclosed pT
         double pt_Rjet = sumPt(near_particles);
         
         // Check if pT is enough
         // Need to have < instead of <= in order to have consistent description with shape variables
         if (pt_Rjet < _ptcut) {
            jets[indices[i]] = NULL;
            continue;
         }

         // If so, select particles in radius
         fastjet::Selector sel_sub = fastjet::SelectorCircle(_Rsub);
         sel_sub.set_reference(my_jets[i]);
         vector<PseudoJet> near_particles_sub = sel_sub(near_particles);
         

         // Calculate enclosed pT
         double pt_Rsub = sumPt(near_particles_sub);
      
         // Need to have < instead of <= in order to have consistent description with shape variables
         if (pt_Rsub/pt_Rjet < _fcut) {
            jets[indices[i]] = NULL;
            continue;
         }
      }
   }

   virtual bool applies_jet_by_jet() const {return false;}

   std::string jetParameterString() const {
      std::stringstream stream;
      stream << "R_jet=" << _Rjet << ", pT_cut=" << _ptcut << ", R_sub=" << _Rsub <<", fcut=" << _fcut;
      return stream.str();
   }

   virtual string description() const {
      return "Shape trimmer, " + jetParameterString();
   }
}; 

//----------------------------------------------------------------------
// Helper for SelectorJetShapeTrimming. 
// Class derived from SelectorWorker: pass, terminator, applies_jet_by_jet, and description are overloaded. 
// This cannot be applied to an individual jet.
class SW_JetShapeTrimming: public SelectorWorker {

private:

   double _Rsub, _fcut;

public:
   
   SW_JetShapeTrimming(double Rsub, double fcut) : _Rsub(Rsub), _fcut(fcut) {};

   virtual bool pass(const PseudoJet &) const {
	if (!applies_jet_by_jet())	
     	throw Error("Cannot apply this selector worker to an individual jet");
      return false;
   }

   virtual void terminator(vector<const PseudoJet *> & jets) const {
      
      // Copy pointers content to a vector of PseudoJets. Check if the pointer has been already nullified. 
      vector<PseudoJet> my_jets;
      vector<unsigned int> indices; 
      for (unsigned int i=0; i < jets.size(); i++){
         if (jets[i]){
            indices.push_back(i);
            my_jets.push_back(*jets[i]);
         } 	
      }
	
      FunctorJetScalarPt sumPt = FunctorJetScalarPt();
      
      // take sum pt of given jet
      double pt_Rjet = sumPt(my_jets);
   
   	for (unsigned int i=0; i < my_jets.size(); i++){

         // Select particles in sub radius
         fastjet::Selector sel_sub = fastjet::SelectorCircle(_Rsub);
         sel_sub.set_reference(my_jets[i]);
         vector<PseudoJet> near_particles_sub = sel_sub(my_jets);
         
         // Calculate enclosed pT
         double pt_Rsub = sumPt(near_particles_sub);
      
         // Need to have < instead of <= in order to have consistent description with shape variables
         if (pt_Rsub/pt_Rjet < _fcut) {
            jets[indices[i]] = NULL;
            continue;
         }
      }
   }

   virtual bool applies_jet_by_jet() const {return false;}

   std::string jetParameterString() const {
      std::stringstream stream;
      stream << "R_sub=" << _Rsub <<", fcut=" << _fcut;
      return stream.str();
   }

   virtual string description() const {
      return "Jet shape trimmer, " + jetParameterString();
   }
}; 


Selector SelectorShapeTrimming(double Rjet, double ptcut, double Rsub, double fcut){
   return Selector(new SW_ShapeTrimming(Rjet,ptcut,Rsub,fcut));
} 

Selector SelectorJetShapeTrimming(double Rsub, double fcut){
   return Selector(new SW_JetShapeTrimming(Rsub,fcut));
} 

//////
//
// Variable pT_cut
//
//////

// helper to sort a vector of vector<double> by comparing just the first entry.
bool _mySortFunction (std::vector<double> v_0, std::vector<double> v_1) { return (v_0[0] > v_1[0]); }

// As described in the appendix of the physics paper, one can construct the inverse 
// of an event shape with respect to pT by calculating all of the possible pTi,R values and sorting them.
// This helper function accomplishes that task.
void JetLikeEventShape_VariablePtCut::_buildStepFunction(const std::vector<PseudoJet> particles) const {
 
   FunctorJetScalarPt sumPt = FunctorJetScalarPt();
   fastjet::Selector sel = fastjet::SelectorCircle(_Rjet);
   _pt_in_Rjet = 0.0;
   _pt_in_Rsub = 0.0;
   
   std::vector< std::vector<double> > myValues;
   myValues.resize(0);

   // this acts much like establishWeights did for JetLikeEventShape,
   // except each value of _pt_in_Rjet is associated with the corresponding
   // weighted measurement.
   for (unsigned int i = 0; i < particles.size(); i++){
      sel.set_reference(particles[i]);
      _nearbyParticles = sel(particles);
      _pt_in_Rjet = sumPt(_nearbyParticles);
      
      if(_trim) {
         assert(_Rsub <= _Rjet);
         fastjet::Selector sel_sub = fastjet::SelectorCircle(_Rsub);
         sel_sub.set_reference(particles[i]);
         std::vector<PseudoJet> near_particles_sub = sel_sub(_nearbyParticles); 
         _pt_in_Rsub = sumPt(near_particles_sub);
         if (_pt_in_Rsub / _pt_in_Rjet >= _fcut) {
            double measurement = _measurement->result(_nearbyParticles);
            std::vector<double> point(2);
            point[0] = _pt_in_Rjet;
            point[1] = measurement*particles[i].pt()/_pt_in_Rjet;
            myValues.push_back(point);
         }
      }

      else {		

         double measurement = _measurement->result(_nearbyParticles);
         std::vector<double> point(2);
         point[0] = _pt_in_Rjet;
         point[1] = measurement*particles[i].pt()/_pt_in_Rjet;
         myValues.push_back(point);
      }
   }



   std::sort (myValues.begin(), myValues.end(), _mySortFunction);

   // Calculating the cumulative sum of the measurements up to a given pt value
   _ptR_values.resize(0);
   _eventShape_values.resize(0);   
   if (!myValues.empty()){
	
      _ptR_values.push_back(myValues[0][0]);  
   	_eventShape_values.push_back(myValues[0][1]);

   	for (unsigned int i = 1; i < myValues.size(); i++){
         _ptR_values.push_back(myValues[i][0]);
         _eventShape_values.push_back(_eventShape_values[i-1] + myValues[i][1]);
      }
   }
   
   myValues.clear();
}

double JetLikeEventShape_VariablePtCut::eventShapeFor(const double ptcut_0) const {

   // TODO: find a more efficient data structure to do this searching

   double eventShape = 0.0;  

   if ( ptcut_0 <=  _ptR_values.back() ) eventShape = _eventShape_values.back();
	
   else if (ptcut_0 >  _ptR_values.back() && ptcut_0 <=  _ptR_values.front()) {
      for (unsigned int i = 0; i < _ptR_values.size()-1; i++){
         if (ptcut_0 <= _ptR_values[i] && ptcut_0 > _ptR_values[i+1]){
            eventShape = _eventShape_values[i];
            break;
         }
      }
   }
   return (eventShape);
}

double JetLikeEventShape_VariablePtCut::ptCutFor(const double eventShape_0) const {

   // TODO: find a more efficient data structure to do this searching

   double ptcut = 0.0;  
   
   double new_eventShape_0 = eventShape_0 - _offset;
   
   if ( new_eventShape_0 < 0 || new_eventShape_0 > _eventShape_values.back() ) {
      cout << ">>> ERROR: "<< _measurement->description() <<" min allowed value is 0, max allowed value is " << _eventShape_values.back() << endl;

   } else if ( new_eventShape_0 <=  _eventShape_values.front() ) ptcut = _ptR_values.front();
   
   else {
      for (unsigned int i = 0; i < _eventShape_values.size()-1; i++){
         if (new_eventShape_0 > _eventShape_values[i] && new_eventShape_0 <= _eventShape_values[i+1]){
            ptcut = _ptR_values[i+1];
            break;
         }
      }
   }
   
   return(ptcut);
}


//////
//
// Njet with variable R_jet
//
//////

// Similar to the same named function in JetLikeEventShape_VariablePtCut, except now we have
// to store a lot more information.
void ShapeJetMultiplicity_VariableR::_buildStepFunction(const std::vector<PseudoJet> particles) const {
   
   unsigned int N = particles.size();

   FunctorJetScalarPt sumPt = FunctorJetScalarPt();
   
   std::vector< std::vector<double> > myValues;
   myValues.resize(0);

   for(unsigned int i = 0; i < N; i++) {
		unsigned int j = i+1;
		while(j < N) {
			vector<double> myPair(3);
			myPair[0] = particles[i].delta_R(particles[j]);
			myPair[1] = i;
			myPair[2] = j;
			myValues.push_back(myPair);
			j++;
		}
   }

   std::sort (myValues.begin(), myValues.end(), _mySortFunction);

   _eventShape_values.resize(0);
   _dR_values.resize(0);

   // Initialize
   double _tot_pt = sumPt(particles);
   if(_tot_pt >= _ptcut) _eventShape_values.push_back(1.0); 
   else _eventShape_values.push_back(0.0);
 
   std::vector<double> pTR(N);
   std::fill(pTR.begin(), pTR.end(), _tot_pt);

   if (_trim) {
      std::vector<double> pTRsub;
      fastjet::Selector sel_sub = fastjet::SelectorCircle(_Rsub);
      for (unsigned int i = 0; i < N; i++) {
         sel_sub.set_reference(particles[i]);
         std::vector<PseudoJet> near_particles_sub = sel_sub(particles); 
         pTRsub.push_back(sumPt(near_particles_sub));
      }
      
      
      for (unsigned int i = 0; i < myValues.size(); i++) {
         
         _dR_values.push_back(myValues[i][0]);
         
         int id1 = (int)myValues[i][1];
         int id2 = (int)myValues[i][2];
         
         pTR[id1] -= particles[id2].pt();
         pTR[id2] -= particles[id1].pt();
         
         double myEventShape = 0;
         
         for(unsigned int j = 0; j < N; j++) {
            if(pTR[j] >= _ptcut && pTRsub[j] / pTR[j] >= _fcut) myEventShape += particles[j].pt() / pTR[j];
         }
         _eventShape_values.push_back(myEventShape);
      }
   }
	
   else {
      
      for (unsigned int i = 0; i < myValues.size(); i++) {
         
         _dR_values.push_back(myValues[i][0]);
         
         int id1 = (int)myValues[i][1];
         int id2 = (int)myValues[i][2];
         
         pTR[id1] -= particles[id2].pt();
         pTR[id2] -= particles[id1].pt();
         
         double myEventShape = 0;
         
         for(unsigned int j = 0; j < N; j++) {
            if(pTR[j] >= _ptcut) myEventShape += particles[j].pt() / pTR[j];
         }
         _eventShape_values.push_back(myEventShape);
      }
   }
   
   _dR_values.push_back(0.0); 
   myValues.clear();

}		

// returns event shape for a given Rjet value
double ShapeJetMultiplicity_VariableR::eventShapeFor(const double Rjet_0) const {

   double eventShape = 0.0;  

   if( Rjet_0 < _Rsub ) cout <<">>> WARNING: R_jet < R_sub= " << _Rsub << endl;

   if( Rjet_0 >= _dR_values.front() ) eventShape = _eventShape_values.front();
   else if ( Rjet_0 < _dR_values.front() && Rjet_0 >= _dR_values.back() ){

     //-GPS very dangerous to do -1 on a size: if size=0 then, because it's unsigned,
     //- you end up with a value that's equal to to 2^nbits-1 where nbits=32 or 64.
   	for(unsigned int i = 0; i < _dR_values.size()-1; i++) {
  		 if( Rjet_0 < _dR_values[i] && Rjet_0 >= _dR_values[i+1]) {
			eventShape = _eventShape_values[i+1];
			break;
  		 }
	}
   }
   else cout <<">>> ERROR: R_jet should be >= 0" << endl;
   return(eventShape);
}


//////
//
// Finding Jet Axes with the Event Shape Density
//
//////

// Find the local axes in a region of size R around each particle
void EventShapeDensity_JetAxes::_find_local_axes(const std::vector<PseudoJet> & particles) const {
   
   _myParticles = particles;
   _N = particles.size();
   
   _axes.resize(0);
   _Njet_weights.resize(0);
   _pt_weights.resize(0);
      
   // Indexing particles (note that this is happening internally, so shouldn't modify user input)
   for (unsigned int i=0; i<_N; i++) _myParticles[i].set_user_index(i);

   // For speed up
   if (_useStorageArray) {
      _myStorageArray.establishStorage(_myParticles,_Rjet, _ptcut);
   }
   
   // Find axis associated with each particle.	
   for (unsigned int i=0; i<_N; i++) {
   
      PseudoJet myParticle = _myParticles[i];
      FunctorJetScalarPt sumPt = FunctorJetScalarPt();

      //Recombiner has to carry user_index information.
      FunctorJetAxis axis(_jetDef); 

      fastjet::Selector sel = SelectorCircle(_Rjet);
      sel.set_reference(myParticle);
      vector<PseudoJet> nearbyParticles;

      if (_useStorageArray) { // start from phi/rapidity strips
         if (!_myStorageArray.aboveCutFor(myParticle)) {
            nearbyParticles.resize(0); //don't do any further analysis if not above jet pTcut
         } else {
            nearbyParticles = sel(_myStorageArray.getStorageFor(myParticle));
         } 
      } else { // use all particles
         nearbyParticles = sel(_myParticles);
      }
      
      int myAxis = -1;
      double pt_w = 0;	
      double Njet_w = 0;	
      
      double pt_in_Rjet = sumPt(nearbyParticles);
      if(pt_in_Rjet >= _ptcut) {
         myAxis = axis(nearbyParticles).user_index();
         pt_w = _myParticles[i].pt(); 	
         Njet_w = pt_w / pt_in_Rjet;
      } 
      
      _axes.push_back(myAxis);
      _pt_weights.push_back(pt_w);
      _Njet_weights.push_back(Njet_w);
      
   }
}

// Take the axes calculated by _find_local_axes and figure out what the final
// global axes will be.  With the global consistency condition, this decreases
// the number of axes under consideration. 
void EventShapeDensity_JetAxes::find_axes_and_weights() const {


   //If requested establish global consistency.
   //Loop over axes and reassign unstable axes until no unstable axis is left.
   if(_applyGlobalConsistency) {
      int n_unstable_axes;
      do {
	n_unstable_axes = 0;	
   	for (unsigned int i=0; i<_N; i++) {
            if(_axes[i]!= -1 && !_isStable(_axes[i])){ 
               n_unstable_axes++;
               _axes[i] = _axes[_axes[i]];
            }	
   	}
	} while(n_unstable_axes>0);
      
   }
 
   //Find distinct axes. 
   //Output a vector of Pseudojets (sorted by pT), with pT of each axis corresponding to the pT weight.
   vector<double> tot_Njet_weights(_N,0), tot_pt_weights(_N,0);
	
   for (unsigned int i=0; i<_N; i++) {
      if (_axes[i] != -1) {
         tot_pt_weights[_axes[i]] += _pt_weights[i];
         tot_Njet_weights[_axes[i]] += _Njet_weights[i];
      }
   }

   _distinctAxes.resize(0);
   _tot_Njet_weights.resize(0);
	
   for (unsigned int i=0; i<_N; i++){
      if(tot_pt_weights[i] > 0) {
         PseudoJet myAxis = _lightLikeVersion(_myParticles[i]) * tot_pt_weights[i];
         _distinctAxes.push_back(myAxis);
      }
   }

   _distinctAxes = sorted_by_pt(_distinctAxes);
   
   //Find the corresponding N_jet weights.
   for(unsigned int i=0; i<_distinctAxes.size(); i++) _tot_Njet_weights.push_back(tot_Njet_weights[_distinctAxes[i].user_index()]);
			
}

// Check stability: a particle defines a stable wta axis if it is its own wta or if it is pointing to a particle that does not pass pTcut.
bool EventShapeDensity_JetAxes::_isStable(const int thisAxis) const {

	bool answer = false;
	if(_axes[thisAxis] == thisAxis || _axes[thisAxis] == -1) answer = true;
	return(answer);
}


} // namespace contrib

FASTJET_END_NAMESPACE