#ifndef __CASUBJET_HH__
#define __CASUBJET_HH__
#include "fastjet/ClusterSequence.hh"
#include <algorithm>
#include <cmath>

// $Id: CASubJet.hh 1727 2010-06-24 11:19:55Z salam $

/// class to help us get a clean (almost parameter-free) handle on
/// substructure inside a C/A jet. It follows the logic described in
/// arXiv:0906.0728 (and is inspired by the original Cambridge
/// algorithm paper in its use of separate angular and dimensionful
/// distances), but provides some extra flexibility.
///
/// It searches for all splittings that pass a symmetry cut (zcut) and
/// then orders them according to some auxiliary scale choice
/// (e.g. jade distance of the splitting, kt distance of the
/// splitting, etc.)
///
/// Code copyright (C) 2009 by Gavin Salam, released under the GPL.
///
/// The code is to be considered BETA.
class CASubJet {
public:

  /// the available choices of auxiliary scale with respect to which
  /// to order the splittings
  enum ScaleChoice {
    kt2_distance, //< usual min(kti^2,ktj^2)DeltaR_{ij}^2
    jade_distance, //< kti . ktj DeltaR_{ij}^2 (LI version of jade)
    jade2_distance, //< kti . ktj DeltaR_{ij}^4 (LI version of jade * DR^2)
    plain_distance, //<           DeltaR_{ij}^2
    mass_drop_distance,  //<    m_jet - max(m_parent1,m_parent2)
    dot_product_distance  //<  parent1.parent2
  };
  /// shortcut name for kt2
  static const ScaleChoice kt_distance = kt2_distance;

  /// constructs and runs;
  ///
  /// you're probably advised to set z_threshold to something non-zero
  ///
  /// kt2_distance may not be the best choice for the scale choice
  ///
  CASubJet(const fastjet::ClusterSequence & cs, 
	   const fastjet::PseudoJet & jet,
	   ScaleChoice scale_choice = kt2_distance,
	   double      z_threshold  = 0
	   );

  /// just constructs
  CASubJet(ScaleChoice scale_choice = kt2_distance,
	   double      z_threshold  = 0);

  /// sets a minimum delta R below which spliting will be ignored
  /// (only relevant if set prior to calling run())
  void set_dr_min(double drmin) {_dr2_min = drmin*drmin;}

  /// If (abs_z_cut) is set to false (the default) then for a
  /// splitting to be considered, each subjet must satisfy 
  /// 
  ///        p_{t,sub} > z_threshold * p_{t,parent}
  ///
  /// whereas if it is set to true, then each subject must satisfy
  ///
  ///        p_{t,sub} > z_threshold * p_{t,original-jet}
  ///
  /// where parent is the immediate parent of the splitting, and
  /// original jet is the one supplied to the run() function.
  ///
  /// Only relevant is called prior to run().
  void set_absolute_z_cut(bool abs_z_cut) {_absolute_z_cut = abs_z_cut;}
  
  /// run the tagger on the given cs/jet
  virtual void run(const fastjet::ClusterSequence & cs, 
		   const fastjet::PseudoJet & jet) {
    aux_ordered.clear();
    _cs  = & cs;
    _jet =   jet;
    _recurse_through_jet(jet);
    sort(aux_ordered.begin(), aux_ordered.end());
  }


  /// class that contains the result
  class JetAux {
  public:
    fastjet::PseudoJet   jet; //< the subjet (immediate parent of splitting)
    double               aux_distance; //< the auxiliary distance between its two subjets
    double               delta_r; //< the angular distance between its two subjets
  };

  /// the vector of jets ordered according to the auxiliary splitting
  /// distance.
  std::vector<JetAux> aux_ordered;

protected:
  void _recurse_through_jet(const fastjet::PseudoJet & jet);
  
  const fastjet::ClusterSequence * _cs;
  fastjet::PseudoJet   _jet;
  ScaleChoice          _scale_choice;
  double               _z_threshold;
  double               _dr2_min;
  bool                 _absolute_z_cut;

};


/// help us to sort through the distances (so largest comes first)
/// [is that a good idea?]
bool operator<(const CASubJet::JetAux & a, const CASubJet::JetAux & b) {
  return a.aux_distance > b.aux_distance;
}


//----------------------------------------------------------------------
inline CASubJet::CASubJet(const fastjet::ClusterSequence & cs, const fastjet::PseudoJet & jet, ScaleChoice scale_choice, double z_threshold) {
  _scale_choice = scale_choice;
  _z_threshold  = z_threshold;
  _dr2_min = 0.0;
  _absolute_z_cut = false;

  run(cs, jet);
}

//----------------------------------------------------------------------
inline CASubJet::CASubJet(ScaleChoice scale_choice, double z_threshold) {
  _scale_choice = scale_choice;
  _z_threshold  = z_threshold;
  _dr2_min = 0.0;
  _absolute_z_cut = false;
}


///----------------------------------------------------------------------
/// work through the jet, establishing a distance at each branching
inline void CASubJet::_recurse_through_jet(const fastjet::PseudoJet & jet) {

  fastjet::PseudoJet parent1, parent2;
  if (_cs->has_parents(jet, parent1, parent2)) {

    /// make sure the objects are not _too_ close together
    if (parent1.squared_distance(parent2) < _dr2_min) return;

    JetAux aux;
    aux.delta_r = sqrt(parent1.squared_distance(parent2));
    // kt distance (squared)
    switch (_scale_choice) {
    case kt2_distance:
      // a standard (LI) kt distance
      aux.aux_distance = parent1.kt_distance(parent2);
      break;
    case jade_distance:
      // something a bit like a mass: pti ptj Delta R_ij^2
      aux.aux_distance = parent1.perp()*parent2.perp()*parent1.squared_distance(parent2);
      break;
    case jade2_distance:
      // something a bit like a mass*deltaR^2: pti ptj Delta R_ij^4
      aux.aux_distance = parent1.perp()*parent2.perp()*pow(parent1.squared_distance(parent2),2);
      break;
    case plain_distance:
      // Delta R_ij^2
      aux.aux_distance = parent1.squared_distance(parent2);
      break;
    case mass_drop_distance:
      // Delta R_ij^2
      aux.aux_distance = jet.m() - std::max(parent1.m(),parent2.m());
      break;
    case dot_product_distance:
      // parent1 . parent2 
      // ( = jet.m2() - parent1.m2() - parent2.m() in a 
      // 4-vector recombination scheme) 
      aux.aux_distance = dot_product(parent1, parent2);
      break;
    default:
      throw fastjet::Error("unrecognized scale choice");
    }
    aux.jet          = jet;
    

    // not very efficient -- sort out later
    if (parent1.perp2() < parent2.perp2()) std::swap(parent1,parent2);

    if (_absolute_z_cut) {
      bool absz1 = parent1.perp() / _jet.perp() >= _z_threshold;
      bool absz2 = parent2.perp() / _jet.perp() >= _z_threshold;
      if (absz1 && absz2) aux_ordered.push_back(aux);
      if (absz1) _recurse_through_jet(parent1);
      if (absz2) _recurse_through_jet(parent2);
    } else {
      double z = parent2.perp()/(parent1.perp()+parent2.perp());
      if (z >= _z_threshold) aux_ordered.push_back(aux);
      _recurse_through_jet(parent1);
      if (z >= _z_threshold) _recurse_through_jet(parent2);
    }
  }
}

#endif // __CASUBJET_HH__