from xboa  import *
from xboa.Hit   import Hit
from xboa.Bunch import Bunch
import sys
import os
import subprocess
from array import array
import math
import operator

import ROOT 

#some input data
filename = sys.argv[1]
filetype = "maus_root_virtual_hit"

print "First loading the data... "
bunch_list = Bunch.new_list_from_read_builtin(filetype, filename)

# now load the ROOT file

root_file = ROOT.TFile(filename, "READ") # pylint: disable = E1101

# and set up the data tree to be filled by ROOT IO

data = ROOT.MAUS.Data() # pylint: disable = E1101
tree = root_file.Get("Spill")
tree.SetBranchAddress("data", data)

bunch_list = Bunch.new_list_from_read_builtin(filetype, filename)

print " no. Virtual Planes: " + str( len(bunch_list) )

emittance4, beta4, Pt = array( 'd' ), array( 'd' ), array('d')
Bz, x, y, z, Pz, Px, Energy, Amplitude = array( 'd' ), array( 'd' ), array( 'd' ), array( 'd' ), array( 'd' ), array( 'd' ), array( 'd'), array( 'd')

Nmu = array( 'd' )


bunch = bunch_list[0]

bunch.root_histogram('pz','MeV/c')
raw_input('Press Enter to exit')


# LOCAL CUT
bunch.cut({'pz':199.}, operator.le)
bunch.root_histogram('pz','MeV/c')

for i in range( len(bunch) ):
    hit = bunch[i]
    print str(i) + " " + str(hit['pz']) +  " local: " + str(hit['local_weight']) + " global: " + str(hit['global_weight'])

# reset weights, to check global cut
bunch.clear_weights()

raw_input('Press Enter to exit')

#GLOBAL CUT 
bunch.cut({'pz':199.}, operator.le, global_cut=True)
bunch.root_histogram('pz','MeV/c')
for i in range( len(bunch) ):
    hit = bunch[i]
    print str(i) + " " + str(hit['pz']) +  " local: " + str(hit['local_weight']) + " global: " + str(hit['global_weight'])
    
raw_input('Press Enter to exit')