#ifndef __JHTOPTAGGER_HH__
#define __JHTOPTAGGER_HH__

// $Id: JHTopTagger.hh 1650 2009-12-17 12:33:35Z salam $

#include "Range.hh"
using namespace std;
using namespace fastjet;

/// An implementation of the Johns Hopkins top tagger. It is based on
/// the method proposed by Kaplan, Rehermann, Schwartz & Tweedie in
/// arXiv:0806.0848. 
///
/// THIS VERSION IS TO BE CONSIDERED BETA.
///
/// Code copyright (C) 2009 by Gavin Salam, released under the GPL.
class JHTopTagger {
public:
  /// Construct the tagger based on a ClusterSequence (assumed to be
  /// C/A based) and a jet, together with the parameters delta_p
  /// (relative momentum cutoff), delta_r (relative angular cutoff)
  /// and the W mass (needed when 
  ///
  JHTopTagger(const fastjet::ClusterSequence & cs, 
	      const fastjet::PseudoJet & jet,
	      double delta_p = 0.10, double delta_r = 0.19,
	      double mW = 81.0);

  /// return the KRST cos(theta_h) variable (see below).
  double cos_theta_h() const {return _cos_theta_h;}

  /// returns true if enough substructure was found for this to 
  /// perhaps be a top
  bool maybe_top() const {return _subjets.size() >= 3;}

  /// returns true if enough substructure was found for this to
  /// perhaps be a top and if W and top jets are in the correct mass
  /// range and the cos theta_h angle is below the max_cos_theta_h
  /// cut.
  bool maybe_top(const Range & top_range, const Range & W_range, double max_cos_theta_h = 1) const {
    return maybe_top() &&
      top_range.contains(_top_subjet.m()) &&
      W_range.contains(_W_subjet.m()) &&
      cos_theta_h() <= max_cos_theta_h
      ;
  }

  /// returns the top subjet
  const PseudoJet & top_subjet() const {return _top_subjet;}

  /// returns the W subjet (not a valid member of the ClusterSequence)
  const PseudoJet & W_subjet() const {return _W_subjet;}
  
  /// returns the full list of subets that were found
  const vector<PseudoJet> & subjets() const {return _subjets;}

  /// routine used internally to carry out one layer of unsplitting
  vector<PseudoJet> split_once(const PseudoJet & startjet);

  /// routine used internally to set the value of cos(theta_h) as described by KRST
  void set_cos_theta_h();


private:
  const ClusterSequence * _cs;
  const PseudoJet *       _jet;
  double _delta_p, _delta_r;
  vector<PseudoJet> _subjets;
  PseudoJet _top_subjet, _W_subjet, _W_subjet1, _W_subjet2;
  double _cos_theta_h;
  double _mW;

};


//----------------------------------------------------------------------
JHTopTagger::JHTopTagger(const fastjet::ClusterSequence & cs, 
			 const fastjet::PseudoJet & jet,
			 double delta_p, double delta_r, double mW) : _cs(&cs), _jet(&jet), 
							   _delta_p(delta_p), _delta_r(delta_r), _mW(mW) {
  
  // do the first splitting
  vector<PseudoJet> split0 = split_once(jet);
  if (split0.size() > 0) {
    _top_subjet = split0[0];
  } else {
    return; // failure...
  }

  // now try a second splitting on each of the resulting objects
  for (unsigned i = 1; i <= 2; i++) {
    vector<PseudoJet> split1 = split_once(split0[i]);
    if (split1.size() > 0) {
      _subjets.push_back(split1[1]);
      _subjets.push_back(split1[2]);
    } else {
      _subjets.push_back(split0[i]);
    }
  }

  
  // make sure things make sense
  if (_subjets.size() < 3) return;

  // now find the pair of objects that is closest 
  // to the W mass
  double dmW_min = 1e200;
  int ii=-1, jj=-1;
  for (unsigned i = 0 ; i < _subjets.size()-1; i++) {
    for (unsigned j = i+1 ; j < _subjets.size(); j++) {
      double dmW = abs(_mW - (_subjets[i]+_subjets[j]).m());
      if (dmW < dmW_min) {
	dmW_min = dmW; ii = i; jj = j;
      }
    }
  }
  // get the two pieces of the W subjet
  _W_subjet1 = _subjets[ii];
  _W_subjet2 = _subjets[jj];
  if (_W_subjet1.perp2() < _W_subjet2.perp2()) swap(_W_subjet1,_W_subjet2);
  // and get the W itself
  _W_subjet = _W_subjet1 + _W_subjet2;


  // and establish the cos(theta_h) angle
  set_cos_theta_h();

}

/// runs the Johns Hopkins decomposition procedure
vector<PseudoJet> JHTopTagger::split_once(const PseudoJet & startjet) {
  PseudoJet this_jet = startjet;
  PseudoJet p1, p2;
  vector<PseudoJet> result;
  while (_cs->has_parents(this_jet, p1, p2)) {
    if (p2.perp2() > p1.perp2()) swap(p1,p2); // order with hardness
    if (p1.perp() < _delta_p * _jet->perp()) break; // harder is too soft wrt original jet
    if ( (abs(p2.rap()-p1.rap()) + abs(p2.delta_phi_to(p1))) < _delta_r) break; // distance is too small
    if (p2.perp() < _delta_p * _jet->perp()) {
      this_jet = p1; // softer is too soft wrt original, so ignore it
      continue; 
    }
    result.push_back(this_jet);
    result.push_back(p1);
    result.push_back(p2);
    break;
  }
  return result;
}



// The helicity angle is a standard observable in top decays,
// used to determine the Lorentz structure of the top-
// W coupling [13]. It is defined as the angle, measured
// in the rest frame of the reconstructed W, between the
// reconstructed top's flight direction and one of the W decay
// products. Normally, it is studied in semi-leptonic top
// decays, where the charge of the lepton uniquely identifies
// these decay products. In hadronic top decays there
// is an ambiguity which we resolve by choosing the lower
// pT subjet, as measured in the lab frame.
void JHTopTagger::set_cos_theta_h() {
  
  // the two jets of interest: top and lower-pt prong of W
  PseudoJet W2  = _W_subjet2;
  PseudoJet top = _top_subjet;
  
  // transform these jets into jets in the rest frame of the W
  W2.unboost(_W_subjet);
  top.unboost(_W_subjet);

  _cos_theta_h = (W2.px()*top.px() + W2.py()*top.py() + W2.pz()*top.pz())/
    sqrt(W2.modp2() * top.modp2());
}


#endif // __JHTOPTAGGER_HH__