#ifndef __CMTOPTAGGER_HH__
#define __CMTOPTAGGER_HH__

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

#include "Range.hh" 
#include "CASubJet.hh"
#include <cmath>

using namespace std;
using namespace fastjet;

/// attempt to implement a simpler alternative to the Johns Hopkins top tagger
///
/// - It looks for the largest splitting (jade*DeltaR^2 distance) consistent 
///   with z > delta_p;
/// 
/// - then it works on the assumption that the contents are the top, and boosts 
///   them into the top rest frame
///
/// - it then runs the (e+e-) kt algorithm so as to get exactly 3 jets and
///   assumes the W is the source of the 2 least energetic ones.
///
/// - Everything is then boosted back into the original frame, where you
///   the user can place their cuts on the top and W masses.
///
/// Note: unlike the JH top tagger, this does not sculpt the W mass
/// distribution too much (i.e. background will not automatically have the
/// W_subjet() peaked at the W mass).  
///
/// Also it does not benefit enormously from a cos theta_h cut (though
/// one can run such a cut in principle, e.g. abs(cos theta_h)<0.65,
/// as well as other cuts on the momentum fractions of the W subjets,
/// in order to adjust S v. B).
///
/// Credits: Gavin Salam, April 2009, unpublished.
///
class CMTopTagger {
public:
  CMTopTagger() {
    _top_subjet.reset(0.0,0.0,0.0,0.0);
    _W_subjet.reset(0.0,0.0,0.0,0.0);
  }

  CMTopTagger(const fastjet::ClusterSequence & cs, 
	      const fastjet::PseudoJet & jet,
	      double zmin = 0.0,
              double mass_max = 0.0);

  vector<PseudoJet> split_once(const PseudoJet & startjet);

  bool maybe_top() const {return _subjets.size() >= 3;}
  //bool maybe_top() const {return _subjets.size() == 3;}
  bool maybe_top(const Range & top_range, const Range & W_range) const {
    return maybe_top() &&
      top_range.contains(_top_subjet.m()) &&
      W_range.contains(_W_subjet.m()) ;
  }

  const PseudoJet & original_jet() const {return *_jet;}

  /// the top_subjet can be legitimately used with the original cluster sequence
  /// (e.g. to get its constituents)
  const PseudoJet & top_subjet() const {return _top_subjet;}

  /// the W subjet can _not_ be legitimately used with the original cs.
  const PseudoJet & W_subjet() const {return _W_subjet;}

  /// Return a vector with the three subjets, numbered such that 
  /// [0]==b 
  /// [1]==W1 (more energetic in top rest frame)
  /// [2]==W2 (less energetic in top rest frame)
  ///
  /// none of these jets can be legitimately used with the original cs
  const vector<PseudoJet> & subjets() const {return _subjets;}

  /// return the momentum fractions contained in the three subjets;
  const vector<double>   & z_values() const {return _z_values;}

  /// [i] is the Delta R separation between the pair that does not contain [i]
  const vector<double>   & dr_values() const {return _dr_values;}

  /// return the z-fraction of the softest constituent
  double zmin() const {return min(min(_z_values[0],_z_values[1]),_z_values[2]);}
  /// cutting on either of the following allows you to adjust S v. B
  /// results (requiring z12max() > 0.1, one gets something quite
  /// similar to the Johns Hopkins top-tagger performance)
  double z12max() const {return max(_z_values[1],_z_values[2]);}
  double z12sum() const {return (_z_values[1] + _z_values[2]);}

  double drmin() const{return min(_dr_values[0],min(_dr_values[1],_dr_values[2]));}
  double drmax() const{return max(_dr_values[0],max(_dr_values[1],_dr_values[2]));}

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

private:
  const ClusterSequence * _cs;
  const PseudoJet *       _jet;
  double _zmin;
  double _mass_max;
  vector<PseudoJet> _subjets;
  vector<double>    _z_values;
  vector<double>    _dr_values;
  PseudoJet _top_subjet, _W_subjet;
  double _cos_theta_h;

  /// sets the value of cos(theta_h) as described by KRST
  void _set_cos_theta_h();


};


//----------------------------------------------------------------------
CMTopTagger::CMTopTagger(const fastjet::ClusterSequence & cs, 
			 const fastjet::PseudoJet & jet,
			 double zmin,
                         double mass_max
                         ) : _cs(&cs), _jet(&jet), 
                                        _zmin(zmin), _mass_max(mass_max), _subjets(0) {


  // get the first sign of major substructure
  CASubJet casub(*_cs, *_jet, CASubJet::jade2_distance,_zmin);

  unsigned isub = 0;
  for (; isub < casub.aux_ordered.size(); isub++) {
    if (casub.aux_ordered[isub].jet.m() < mass_max || mass_max == 0.0) break;
  }
  if (isub >= casub.aux_ordered.size()) return;
  
  // now try looking at constituents
  _top_subjet = casub.aux_ordered[isub].jet;
  vector<PseudoJet> constituents = _cs->constituents(_top_subjet);
  if (constituents.size() < 3) return;

  // consider them in the top frame
  for (unsigned i = 0; i < constituents.size(); i++) {
    constituents[i].unboost(_top_subjet);
  }

  // recluster them with e+e- kt alg. casting the event to 3jets
  ClusterSequence   csee(constituents,JetDefinition(ee_kt_algorithm));
  _subjets = sorted_by_E(csee.exclusive_jets(3));

  // now go back to lab frame
  for (unsigned i = 0;  i < 3; i++) {_subjets[i].boost(_top_subjet);}

  // least two energetic in the e+e- frame give the W
  _W_subjet = _subjets[1] + _subjets[2];

  // get z-values according to e+e- frame numbering
  _z_values.resize(3);
  for (unsigned i = 0; i < 3; i++) {_z_values[i] = _subjets[i].perp()/_top_subjet.perp();}

  // also get dr values
  _dr_values.resize(3);
  _dr_values[0] = sqrt(_subjets[1].squared_distance(_subjets[2]));
  _dr_values[1] = sqrt(_subjets[0].squared_distance(_subjets[2]));
  _dr_values[2] = sqrt(_subjets[0].squared_distance(_subjets[1]));

  // for what it's worth
  _set_cos_theta_h();
}


//
// The following text is taken from the Johns Hopkins paper
//
// 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.
//
// NB: this is not of particular use for this class, but we've
//     left it in anyway.
void CMTopTagger::_set_cos_theta_h() {
  
  // the two jets of interest: top and lower-pt prong of W
  PseudoJet W2  = _subjets[1].perp2() < _subjets[2].perp2() ? _subjets[1] : _subjets[2];
  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 // __CMTOPTAGGER_HH__