#######################################################################################
# Functions demonstrating the use of the TrackNav, TrackCursor and TrackNode classes  
# All require a ROOT file generated with tracking information ON.                     
#
# see https:##github.com#snoplus#rat#blob#master#example#macros#tracking#tracking.mac 
#     http:##snopl.us#docs#rat#user_manual#html#node214.html
#                                                                                    
# Particle tracks are divided into a series of backward looking 'Nodes' containing   
# the step information. The TrackCursor class controls the navigation through these  
# nodes and between tracks.                                                          
#                                                                                    
# Secondary particles appear new 'child' tracks at the node that created them.  
#                                                                                    
# Primary particles are the children of an inital track with a single node
# representing the event 
#                                                                                  
#  
# J dunger - 2014-08-20 <jack.dunger@ox.ac.uk> : New file   
#######################################################################################

import ROOT
import rat

# A function to plot the primary particle momenta for a set of events
def plot_primary_mom (filename):
    hist   = ROOT.TH1D("","",10 ,0,5)
    # loop over events
    for ds,run in rat.dsreader(filename):
        nav = ROOT.RAT.TrackNav(ds) 
        cursor   = nav.Cursor()                     # cursor at head node of track
        # head node is a track of its own
        # real tracks for generated particles are its children 

        for child in range(0,cursor.ChildCount()):
            node = cursor.Child(child)              # grab node at start of child track
            hist.Fill(node.GetMomentum().Mag())     # nodes contain particle information
            print node.GetParticleName() 


    return hist
            

# A function to plot the 1st primary particle's displacement
def plot_primary1_disp(filename):
    hist   = ROOT.TH1D("","",10 ,0,1000)
    # event loop
    for ds,run in rat.dsreader(filename):
        nav = ROOT.RAT.TrackNav(ds) 
        cursor   = nav.Cursor()      
        
        if(cursor.ChildCount()!=0):
            cursor.GoChild(2)              # move cursor to start of child track
            node     = cursor.Here()       # grab the node here
            init_pos = node.GetPosition()
            
            node = cursor.GoTrackEnd()     # returns correct node AND moves cursor
            fin_pos  = node.GetPosition()
            
            disp = (fin_pos - init_pos).Mag()
            hist.Fill(disp)

    return hist


# A function to find the processes with secondaries
def find_secondaries(filename):
    # event loop
    for ds,run in rat.dsreader(filename):
        nav = ROOT.RAT.TrackNav(ds) 
        cursor   = nav.Cursor()      
        
        # primary particle loop
        for child in range(0,cursor.ChildCount()):
            cursor.GoChild(child)

            #Step along particle track
            while(cursor.IsTrackEnd()==False):
                node = cursor.Here()
                if (cursor.ChildCount() != 0):
                    print node.GetProcess()
                cursor.GoNext()

            cursor.GoTrackStart()
            cursor.GoParent()

    return