#!/usr/bin/env python

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

 Produce summary file in offline format that has hits removed, and only the
 best tracks for each reconstruction algorihtm. 

 options:
 -f : input file                                     default=
 -o : output file                                    default=dst.root
 -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  
 -r : ids of reco algorithms                         default=4000 101
 -H : print header of output file                    default=False
 -s : sample/prescale: process every nth event       default=1
 -n : process this many input events (before prescale) default = -1
"""


import sys, aa
from copy import copy
from ROOT import EventFile, best_reco_track, Trk
import ROOT
options = aa.Options( __doc__, sys.argv[1:] ) 



infiles, outfile = options.f.strip() , options.o

# if multiple input files, put them in list
if " " in infiles : infiles = infiles.split()

# make list of rectypes to put in output file
rec_types = list( map( int , options.r.split() ) )


f = EventFile( infiles )

if options.H : print ( f.header )

# this helps speed things up when reading tree already in aanet format
#f.roottree().SetBranchStatus("mc_hits*",False)
#f.roottree().SetBranchStatus("hits*",False)

# If we set an output file, the output TTree is filled automatically
f.set_output( outfile )    # this opens an aanet output file
f.autosave = False

# Add another tree as "fried"

curdir = ROOT.gDirectory
f.output_file.cd()
T = ROOT.TNtuple("T","my additial variables", "nhits:ntrighits");
f.output_tree.AddFriend(T)
curdir.cd()


# write the header as TNamed's

f.output_file.mkdir("HeadDir")
header_dir = f.output_file.GetDirectory("HeadDir")
if not header_dir : 
	print ("failed to get header dir ")
	sys.exit()

for k,v in f.header :
	header_dir.WriteTObject( ROOT.TNamed( k,v )  )



# reading is like this:
keys= f.output_file.GetDirectory("HeadDir").GetListOfKeys()
print ('---------------------------------------------------------')
print( "\n".join( k.GetName()+": "+k.GetTitle() for k in keys ) ) 
print ('---------------------------------------------------------')


# add macro to read header fields
# in root do init->Exec() and then use get_header
code = """
string get_header( string key ) {
	auto d = gFile->GetDirectory("HeadDir");
	return d->Get( key.c_str() )->GetTitle();
}

void init() {
	cout << " hello. you can now use get_header( key ) to read header information " << endl;
}
"""
macro =  ROOT.TMacro()
for l in code.splitlines():	macro.AddLine( l );
macro.SetName("init")
f.output_file.WriteTObject( macro );



for evt in f : 

	# prescale (use global index to correctly prescale many small input files)
	if f.global_index() % options.s >  0 : continue

	if options.n > 0 and f.global_index() >= options.n : 
		print("requested number of events (",options.n,") reached")
		break;

	if options.p : print( f.global_index() , f.index , evt)

	# record summary information / derived quantities in out friend Tree


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

	# this clears the hits and mc hits and removes the hit-refs from the tracks
	evt.clear_hits();


	# keep only the best reconstruction tracks
	keeptrks=[]
	for rec_type in rec_types :
		t =	best_reco_track( evt , rec_type )     # might return nullptr
		keeptrks.append( Trk(t) if t else Trk() ) # ..in which case, store emtpy 

	#print(f.current_filename() , rec_types, keeptrks)
	evt.trks.clear();
	for t in keeptrks : evt.trks.push_back( t )

	#for the moment, keep only the first two MC tracks (nu + muon )
	evt.mc_trks.resize(2);

	f.write()

del f # this is needed in interactive mode, it seesm :(
#