// $Id$
//
// Copyright (c) 2016,
// Gavin Salam, Gregory Soyez, Jesse Thaler
//
//----------------------------------------------------------------------
// 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/>.
//----------------------------------------------------------------------

#ifndef __FASTJET_CONTRIB_CARTESIANJET_HH__
#define __FASTJET_CONTRIB_CARTESIANJET_HH__

#include <fastjet/internal/base.hh>
#include <fastjet/JetDefinition.hh>
#include <fastjet/PseudoJet.hh>
#include <fastjet/Error.hh>
#include <vector>

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

namespace contrib {
  
  //------------------------------------------------------------------------
  /// \class CartesianCoordinate
  /// This class stores the x, y coordinates and optionally, the intensity of an entry
  /// TODO:  Allow the inclusion of arbitrary user information (templated? allow inheritance?)
  class CartesianCoordinate : public PseudoJet::UserInfoBase {
    
  public:
    // empty constructor for zero entry
    CartesianCoordinate()
    : _x(0), _y(0), _intensity(0), _index(0)
    {}
    
    // default constructor
    CartesianCoordinate(double x, double y, double intensity = 1.0, int index = 0)
    : _x(x), _y(y), _intensity(intensity), _index(index)
    {}
    
    // accessing the coordinates, intensity, and optional user index
    double x() const {return _x;}
    double y() const {return _y;}
    double intensity() const {return _intensity;}
    int user_index() const {return _index;}
    
  private:
    double _x;          // x coordinate
    double _y;          // y coordinate
    double _intensity;  // intensity
    int    _index;      // user index
  };
  
  
  //------------------------------------------------------------------------
  /// \class CartesianJet
  /// Inherits from PseudoJet and contains a pointer to a CartesianCoordinate.
  /// Pointer is needed to use existing user_info structure of PseudoJet
  class CartesianJet : public PseudoJet {
    
  public:
    // empty constructor
    CartesianJet()
    : PseudoJet(0,0,0,0) {
      set_user_info(new CartesianCoordinate()); // create dummy coordinate
    }
    
    // copy constructor has to be non-trivial to make new copy of CartesianCoordinate
    CartesianJet(const CartesianJet & cj)
    : PseudoJet(cj) {
      set_user_info(new CartesianCoordinate(cj.coordinate()));
    }
    
    // Mapping a PseudoJet into a CartesianJet, assuming is has CartesianCoordinate
    // as user_info
    CartesianJet(const PseudoJet & pj)
    : PseudoJet(pj) {
      if ( !has_user_info<CartesianCoordinate>() ) {
        throw Error("CertesianJet:  Constructor expects PseudoJet to have CartesianCoordinate as its user_info().");
      }
      set_user_info(new CartesianCoordinate(pj.user_info<CartesianCoordinate>()));
    }
    
    // getting access to underlying coordinate (returns const reference)
    const CartesianCoordinate & coordinate() const { return user_info<CartesianCoordinate>();}

    // accessing coordinates and intensities (off-loaded to CartesianCoordinate)
    double x() const {return coordinate().x();}
    double y() const {return coordinate().y();}
    double intensity() const {return coordinate().intensity();}
    
    // GPS: as an interim alternative to user info, it might be useful to
    // add a user_index?
    // JDT:  I'm trying to hijack user index of Pseudojet, not sure if this is best
    int user_index() const {return coordinate().user_index();}
    
    std::vector<CartesianJet> cartesian_constituents() const {
      std::vector<PseudoJet> my_constituents = constituents(); // from PseudoJet base class
      return std::vector<CartesianJet>(my_constituents.begin(), my_constituents.end());
    }
    
  protected:
    
    //CartesianRange is responsible for setting the proper pseudojet
    friend class CartesianRange;
    
    // Turning a CartesianCoordinate into a CartesianJet.
    // The underlying Pseudojet starts off as zero and
    // has to be set, e.g., by CartesianRange
    CartesianJet(const CartesianCoordinate & coord)
    : PseudoJet(0,0,0,0) {
      set_user_info(new CartesianCoordinate(coord));
    }
    
  private:
    // making private so user cannot change (it comes from CartesianCoordinate
    void 	set_user_index(const int index);
    
  };
  
  
  //------------------------------------------------------------------------
  /// \class CartesianRange
  /// This class converts CartesianCoordinates into PseudoJets, assuming a given
  /// x and y range
  class CartesianRange {
  public:
    
    static const double max_azimuthal_value; //= M_PI/2.0;
    
    /// default constructor
    CartesianRange(double xmin, double xmax, double ymin, double ymax)
    : _xmin(xmin), _xmax(xmax), _ymin(ymin), _ymax(ymax)
    {
      //rescale factor to put angles between 0 and max_azimuthal_value
      _scale_factor = (_ymax - _ymin)/ max_azimuthal_value;
    }
    
    // convert from CartesianCoordinate to CaretsianJet
    CartesianJet create_cartesian_jet(const CartesianCoordinate& coord) const {
      
      // create CartesianJet with zero underlying PseudoJet, to be set later
      CartesianJet cj(coord);
      
      // get CartesianCoordinate information
      double x = coord.x();
      double y = coord.y();
      double intensity = coord.intensity();
      
      // GPS: I would perhaps make these "assert", because something that fails here
      // is a sign of an internal inconsistency
      // JDT: Not necessarily, if the user makes their own CartesianRange
      
      // issue warnings if range is wrong (might still work)
      if (x < _xmin) _out_of_range_warning.warn("CartesianRange warning: CartesianCoordinate has x < xmin.");
      if (x > _xmax) _out_of_range_warning.warn("CartesianRange warning: CartesianCoordinate has x > xmax.");
      if (y < _ymin) _out_of_range_warning.warn("CartesianRange warning: CartesianCoordinate has y < ymin.");
      if (y > _ymax) _out_of_range_warning.warn("CartesianRange warning: CartesianCoordinate has y > ymax.");
      
      // rescale to rapidity and azimuthal values
      double rap = (x - _xmin) / _scale_factor;
      double phi = (y - _ymin) / _scale_factor;
      
      // Define underlying PseudoJet
      // Note that you have to call reset_momentum_PtYPhiM, not reset_PtYPhiM
      // to avoid reseting user information
      cj.reset_momentum_PtYPhiM(intensity,rap,phi);
    
      return cj;
    }
    
    // convert vector<CartesianCoordinate> to vector<CartesianJet>
    std::vector<CartesianJet> create_cartesian_jets(const std::vector<CartesianCoordinate> & coords) const {
      std::vector<CartesianJet> cjs;
      for (int i = 0; i < coords.size(); i++) {
        cjs.push_back(create_cartesian_jet(coords[i]));
      }
      return cjs;
    }
    
    // get scale factor
    double scale_factor() const {return _scale_factor;}
    
    // get rescaled jet radius value
    double pseudojet_radius(double cj_radius) const {
      double radius = cj_radius/_scale_factor;
      if (radius > max_azimuthal_value) {
        // GPS: again this should be assert (possibly with tolerance?) because
        //      if this happens, it's a sign of an internal inconsistency
        _too_large_radius.warn("CartesianRange warning: resulting radius is beyond max_azimuthal_value; wrapping may occur");
      }
      
      // TODO:  Should we rescale to make it possible to use bigger radius values?
      return radius;
    }
    
  private:
    // minimum and maximum cartesian ranges
    double _xmin;
    double _xmax;
    double _ymin;
    double _ymax;
    
    // calcuated scale factors
    double _scale_factor;
    
    // various warnings
    mutable LimitedWarning _out_of_range_warning;
    mutable LimitedWarning _no_cartesianjet;
    mutable LimitedWarning _too_large_radius;
    
  };
  
  //------------------------------------------------------------------------
  /// \class CartesianRecombiner
  /// This is a custom PseudoJet recombiner the not only adds the Cartesian
  /// jet information but also uses pt-scheme recombination
  class CartesianRecombiner : public JetDefinition::Recombiner {
    
  public:
    // Constructor with delta (effectively pt^delta-scheme)
    CartesianRecombiner(double delta = 1.0)
    : _delta(delta)
    {}
    
  private:
    double _delta;
    
    virtual std::string description() const {
      std::ostringstream ostr;
      ostr << "CartesianRecombiner (delta = " << _delta << ")";
      return ostr.str();
    }
    
    // This combiner information
    virtual void recombine(const fastjet::PseudoJet & pa,
                           const fastjet::PseudoJet & pb,
                           fastjet::PseudoJet & pab) const {
      
      
      double perp_a = pa.perp();
      double perp_b = pb.perp();
      double perp_ab = perp_a + perp_b;
      
      // setting initial values (used also if equal pTs)
      double ratio    = 1.0;
      double weight_a = 0.5;
      double weight_b = 0.5;
      
      if (perp_a > perp_b){
        ratio = pow(perp_b/perp_a,_delta);
        weight_a = 1.0  /(1.0 + ratio);
        weight_b = ratio/(1.0 + ratio);
      } else if (perp_b > perp_a){
        ratio = pow(perp_a/perp_b,_delta);
        weight_b = 1.0  /(1.0 + ratio);
        weight_a = ratio/(1.0 + ratio);
      }
      
      if (perp_ab != 0.0) {
        double y_ab = (weight_a * pa.rap() + weight_b * pb.rap());
        
        double phi_a = pa.phi(), phi_b = pb.phi();
        if (phi_a - phi_b > pi)  phi_b += twopi;
        if (phi_a - phi_b < -pi) phi_b -= twopi;
        double phi_ab = (weight_a * phi_a + weight_b * phi_b);
        
        pab.reset_PtYPhiM(perp_ab, y_ab, phi_ab);
        
        // Do the same for underlying CartesianJet
        
        if (pa.has_user_info<CartesianCoordinate>() && pb.has_user_info<CartesianCoordinate>()) {

          CartesianCoordinate cja = pa.user_info<CartesianCoordinate>();
          CartesianCoordinate cjb = pb.user_info<CartesianCoordinate>();
          
          double x = weight_a * cja.x() + weight_b * cjb.x();
          double y = weight_a * cja.y() + weight_b * cjb.y();
          double intensity = cja.intensity() + cjb.intensity();
          
          CartesianCoordinate * new_cj = new CartesianCoordinate(x,y,intensity);
          pab.set_user_info(new_cj);
        }
      }
      else {
        pab.reset(0.0,0.0,0.0,0.0);
        pab.set_user_info(new CartesianCoordinate());
      }
    }
    
  };
  
  //------------------------------------------------------------------------
  /// \class CartesianJetDefinition
  /// This wraps the gen-kt algorithm, allowing for pt^delta scheme, but no other options
  class CartesianJetDefinition {
    
  public:
    // Constructor, default is anti-kT with pT-scheme recombination
    CartesianJetDefinition(double cartesian_radius, double p = -1.0, double delta = 1.0);
    
    // running the jet algorithm
    std::vector<CartesianJet> operator() (const std::vector<CartesianCoordinate> & coord_input) const {
      
      // GPS: 0-size vector is a special case that should be treated smoothly;
      if (coord_input.size() == 0) return std::vector<CartesianJet>();
      
      // First figure out size of image
      double xmin = std::numeric_limits<double>().max();
      double xmax = std::numeric_limits<double>().min();
      double ymin = std::numeric_limits<double>().max();
      double ymax = std::numeric_limits<double>().min();
      for (int i = 0; i < coord_input.size(); i++) {
        // GPS: replace these with std min/max, which is more idiomatic
        xmin = std::min(xmin,coord_input[i].x());
        xmax = std::max(xmax,coord_input[i].x());
        ymin = std::min(ymin,coord_input[i].y());
        ymax = std::max(ymax,coord_input[i].y());
      }
      
      // GPS: this is necessary to avoid problems that can occur with
      //      a limited yrange, namely the fact that the cartesian
      //      radius can end up translating to a large FJ radius,
      //      which then causes problems with the periodicity
      ymax = std::max(ymax, ymin + _cartesian_radius);
      
      // Next figure out rescaling of jet radius.
      CartesianRange range(xmin,xmax,ymin,ymax);
      double radius = range.pseudojet_radius(_cartesian_radius);
      
      // Define jet algorithm, crucially with recombiner
      JetDefinition jet_def(genkt_algorithm,radius, _p);
      jet_def.set_recombiner(&_recombiner);
      
      // Convert CartesianCoordinates to CartesianJets
      std::vector<CartesianJet> cjs_input = range.create_cartesian_jets(coord_input);
      
      // Run clustering
      ClusterSequence * clust_seq = new ClusterSequence(cjs_input,jet_def);
      
      // Get inclusive jet output
      std::vector<PseudoJet> pjs_output = sorted_by_pt(clust_seq->inclusive_jets(0.0));
      
      // Convert to CartesianJets
      std::vector<CartesianJet> cjs_output(pjs_output.begin(),pjs_output.end());

      //Tell cluster sequence to delete self
      clust_seq->delete_self_when_unused();
      
      return cjs_output;
    }
    
  private:
    double _cartesian_radius;
    double _p;
    CartesianRecombiner _recombiner;
    
  };
  
} // namespace contrib

FASTJET_END_NAMESPACE

#endif  // __FASTJET_CONTRIB_CartesianJet_HH__