#!/usr/bin/env python
# $Id: HWGTopExample.py 873 2010-07-07 00:04:39Z verm $

__author__ = 'Christopher K. Vermilion'
__email__ = 'verm@uw.edu'
__modified__ = '$Date: 2010-07-06 17:04:39 -0700 (Tue, 06 Jul 2010) $'
__copyright__ = 'Copyright 2010 Christopher K. Vermilion'
__license__ = 'GPL http://www.gnu.org/licenses/gpl.html'

import sys
from  SpartyJetConfig import *
from ROOT import gSystem

gSystem.Load('libFastPruneTool.so')
#===============================================

if len(sys.argv) < 3:
  print 'Usage: ./HWGTopExample.py input_file output_file'

# See the SpartyJet example examples_py/TopTaggerExample.py for more details
#  on the implementation of these taggers.
#
# In taking jets from the output of the taggers, the top jet is taken to be
#  the sum of the subjets(), not top_subjet().  The mass of the W subjet is
#  stored as "mW", and the value of cos_theta_h() is stored as cosThetaH.  To
#  add cuts on these, use the JetMomentSelectorTool, an example of which is
#  commented out below under the JH tagger setup.


# Create a jet builder---------------------------
builder = SJ.JetBuilder(SJ.ERROR)

# Configure input -------------------------------

input = SJ.HuskyInput(sys.argv[1])
builder.configure_input(input)

# Configure algorithms --------------------------

# Run AntiKt10, then fork off to compare against top taggers
AntiKt10 = SJ.FastJet.FastJetFinder('AntiKt10',fj.antikt_algorithm,1.0,False)
builder.add_default_alg(AntiKt10)
# cut on pT *before* forking to save time
builder.add_jetTool(SJ.JetPtSelectorTool(400*input.getGeV()), 'AntiKt10')
AKTparent = SJ.ForkToolParent('AntiKt10Parent')
builder.add_jetTool(AKTparent, 'AntiKt10')

# recluster with CA, then fork again for different taggers
AKTchild = SJ.ForkToolChild(AKTparent, 'CA10')
builder.add_custom_alg(SJ.JetAlgorithm(AKTchild))
builder.add_jetTool(SJ.FastJet.FastJetRecluster('CA10cluster', fj.cambridge_algorithm, 1.5, False), 'CA10')
CAparent = SJ.ForkToolParent('CA10Parent')
builder.add_jetTool(CAparent, 'CA10')

# JH Top Tagger
child = SJ.ForkToolChild(CAparent, 'CA10JH')
builder.add_custom_alg(SJ.JetAlgorithm(child))

# CM Top Tagger
child2 = SJ.ForkToolChild(CAparent, 'CA10CM')
builder.add_custom_alg(SJ.JetAlgorithm(child2))

# Pruning (not technically a "top tagger", but useful for finding tops)
child4 = SJ.ForkToolChild(CAparent, 'CA10Pruned')
builder.add_custom_alg(SJ.JetAlgorithm(child4))

# ------------------------------------------------

# Input min pt cut
#builder.add_jetTool_input(SJ.JetPtSelectorTool(0.5*input.getGeV()))
# PDG id cut -- remove muons (13) and neutrinos (12, 14, 16)
ids = stdVector(-12, 12, 14, -14, 16, -16, 13, -13)
builder.add_jetTool_input(SJ.JetInputPdgIdSelectorTool(ids))
# Input eta cut
builder.add_jetTool_input(SJ.JetEtaCentralSelectorTool(-5.0,5.0))

#======Here is where the algorithms defined above will run ======

# Set up JH top tagger
JHtagger = fj.JHTopTagger(0.1, 0.19, 81.0)
JHtool = SJ.FastJet.TopTaggerTool(fj.JHTopTagger)(JHtagger)
builder.add_jetTool(JHtool, 'CA10JH')
# add a cut on cos_theta_h, which is stored by the tagger
#builder.add_jetTool(SJ.JetMomentSelectorTool(float)('cosThetaH', -0.65, 0.65), 'CA10JH')

# Set up CM top tagger
CMtagger = fj.CMTopTagger(0.1, 0.0)
CMtool = SJ.FastJet.TopTaggerTool(fj.CMTopTagger)(CMtagger)
builder.add_jetTool(CMtool, 'CA10CM')

# Set up pruning
big_CA_def = fj.JetDefinition(fj.cambridge_algorithm, 3.14*0.5)
CAprune_tool = SJ.FastJet.FastPruneTool(big_CA_def)
builder.add_jetTool(CAprune_tool, 'CA10Pruned')

# Measure heavier subjet mass (assumed to be W mass for pruned jets)
subjetMoment = SJ.HeavierSubjetMass("subjetM")
builder.add_jetTool(SJ.JetMomentTool("SubjetMassTool",subjetMoment))

# Configure output--------------------------------
builder.add_text_output("../data/output/HWGTopExample.dat")
builder.configure_output("SpartyJet_Tree",sys.argv[2])

# Run SpartyJet
builder.print_event_every(100)
builder.process_events()