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)

vhit_z = ROOT.TH1D("z[m]",  #pylint: disable = E1101

                                  " z [m]; ",

                                  50, -6, 6)

vhit_Pz = ROOT.TH1D("Pz[MeV/c]",  #pylint: disable = E1101

                                   " Pz [MeV/c]; ",

                                   50, 170, 210)

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_list[100].cut({'y':-300.}, operator.le, global_cut=True)

bunch_list[100].cut({'y':300.}, operator.ge, global_cut=True)

bunch_list[100].root_histogram('y','mm')

raw_input('Press Enter to exit')

bunch_list[100].clear_weights()

bunch_list[100].cut({'y':-300.}, operator.le, global_cut=False)

bunch_list[100].cut({'y':300.}, operator.ge, global_cut=False)

bunch_list[100].root_histogram('y','mm')

raw_input('Press Enter to exit')

for Plane in range( len(bunch_list) ):

    #    bunch_list[Plane].cut({'px':0.})#, operator.le, global_cut=True)

    #print str(Plane) + " loop 2"

    my_bunch = bunch_list[Plane]

    emittance4.append( my_bunch.get_emittance(['x','y'])  )

    beta4.append( my_bunch.get_beta(['x','y']) /10 )

    x.append(  my_bunch.get('mean', ['x']  ))

    y.append(  my_bunch.get('mean', ['y']  ))

    z.append(  my_bunch.get('mean', ['z']  ))

    Pz.append( my_bunch.get('mean',['pz']  ))

    Px.append( my_bunch.get('mean',['px']  ))

    Energy.append(  my_bunch.get('mean', ['energy']  ) ) 

    Pt.append( (  (my_bunch.get( 'mean',['px'] )  )**2 +  (my_bunch.get( 'mean',['py']) )**2   ) ** 0.5 )

    Bz.append(my_bunch.get('mean',['bz']) * 1000 )

    #Amplitude.append(my_bunch.get('mean',['amplitude x y'] ))

    Nmu.append( my_bunch.bunch_weight() )

cX = ROOT.TCanvas("cX", # pylint: disable = E1101

                        "cX")

cX.Divide(2)

cX.cd(1)                      

g27 =  ROOT.TGraph(len(z), z, Nmu)

g27.Draw("AL")

cX.cd(2)  

g28 =  ROOT.TGraph(len(z), z, x )

g28.Draw("AL")

c2 = ROOT.TCanvas("c2", # pylint: disable = E1101

                        "c2")

c2.Divide(2)

c2.cd(1)

g21 =  ROOT.TGraph(len(z), z, emittance4)

g21.SetTitle("#epsilon 4d [mm]")

g21.GetXaxis().SetTitle("z [m]")

g21.GetYaxis().SetTitle("#e [mm]")

g21.Draw("AL")

c2.cd(2)

g22 = ROOT.TGraph(len(z), z, Energy)

g22.SetTitle("Energy [MeV]")

g22.GetXaxis().SetTitle("z [m]")

g22.GetYaxis().SetTitle("Energy [MeV]")

g22.Draw("AL")

c3 = ROOT.TCanvas("c3", # pylint: disable = E1101

                        "c3")

c3.Divide(2)

c3.cd(1)

g31 = ROOT.TGraph(len(z), z, Pt)

g31.SetTitle("Pt [MeV/c]")

g31.GetXaxis().SetTitle("z [m]")

g31.GetYaxis().SetTitle("")

g31.Draw("AL")

c3.cd(2)

g32 = ROOT.TGraph(len(z), z, Bz)

g32.SetTitle("Bz [T]")

g32.GetXaxis().SetTitle("z [m]")

g32.GetYaxis().SetTitle("Bz [T]")

g32.Draw("AL")

c1 = ROOT.TCanvas("c1", # pylint: disable = E1101

                        "c1")

c1.Divide(2)

c1.cd(1)

g = ROOT.TGraph(len(z), z, Pz)

g.SetTitle("Pz [MeV/c]")

g.GetXaxis().SetTitle("z [m]")

#g.GetYaxis().SetTitle("Pz [MeV/c]")

g.Draw("AL")

c1.cd(2)

g1 = ROOT.TGraph(len(z), z, beta4)

g1.SetTitle("#beta [cm]")

g1.GetXaxis().SetTitle("z [m]")

#g1.GetYaxis().SetTitle("#beta [cm]")

g1.Draw("AL")

print " change in emittance: " + str( abs( emittance4[0] - emittance4[ len(emittance4) - 1 ]) )

raw_input('Press Enter to exit')

raise SystemExit