#!/usr/bin/env python

"""
 friend_ntuple.py any_to_aa.py -f input_file -o output_file

 This example illustrates how to use ROOT Tree-friend
 mechnism to 'add' data to the Evt's. This way, a (flat) ntupel
 can be made with user-derived data, while maintaining the
 link to the Evt.

 options:
 -f : input file default=/sps/km3net/users/heijboer/example_files/JGandalf_aanet_muon-CC_3-100GeV_1.root
 -o : output file                                    default=friend.aa.root
 -n : process this many events,                      default=10
 -p : print 1-line summaries of the events           default=True
 -h : this help message and exit                     default=False       
 -i : Python interative mode (prompt when done)      default=False  
"""

import sys, aa
from ROOT import *
options = aa.Options( __doc__, sys.argv[1:] ) 



f = EventFile( options.f )

f.set_output( options.o ) # Creates a TTree (f.output_tree) to store the aanet evts.

f.autosave = False        # we are not going to keep all events,
                          # but call f.output_tree() ourselves.

# Add another tree as "fried"
T = TNtuple("T","my additial variables", "nhits:ntrighits:has_many_trighits");
f.output_tree.AddFriend(T)


for evt in f : 

    if f.index >= options.n : break;

    if options.p : print evt
    
    # demand neutrino energy > 10 GeV
    nu = evt.mc_trks[0]
    if nu.E<10 : continue

    nhits     = evt.hits.size()
    ntrighits = sum( h.trig for h in evt.hits )
    hasmanytrg= ntrighits > 20
    T.Fill( nhits, ntrighits, hasmanytrg );

    # Now, we erase the hits from the output file to save space.
    evt.hits.clear()
    
    # If we get here, we "fill" the output tree 
    # This will also fill the friend tree, T
    f.fill_output_tree()

del f

# Now let's check. Using TTree::Draw and TTree::Scan, the 
# friends variables seems to live in the same TTree.

ff = TFile( options.o );
ff.ls()
E.Scan("nhits:log10(mc_trks[0].E)") # nhits vs Enu

# note E just exists since we opened the file 
# (this just like in a CINT session).