### Aufgabenstellung
# a) Öffnet die Datei A4_file.root mit dem TBrowser und verschafft euch einen
#    Überblick über die enthaltenen Variablen und deren Bedeutung.
# b) Berechnet die invariante Masse des 𝐽⁄𝜓-Mesons (relativistisch) und fügt diese
#    als neue Variable dem Tree hinzu.
# c) Schaut euch die berechnete Massenverteilung des 𝐽⁄𝜓-Mesons an.
# d) Berechnet die Effizienz der Selektion 2400 MeV/c2 < m( 𝐽⁄𝜓) < 3300 MeV/c2 und
#    berechnet den Fehler als Binomialfehler.

import ROOT as R
import numpy as np

### Aufgabenteil a) und c)
# Terminal:             root path/to/file.root
# In der ROOT-Shell:    new TBrowser()

### Aufgabenteil b)
my_file = R.TFile('/ceph/programmierkurs/root/A4_file.root', 'READ')
my_tree = my_file.Get('DecayTuple')

new_file = R.TFile('A4_file_Jpsi.root', 'RECREATE')
new_tree = my_tree.CloneTree()

E1_PE = np.zeros(1, dtype=np.float64)
E1_PX = np.zeros(1, dtype=np.float64)
E1_PY = np.zeros(1, dtype=np.float64)
E1_PZ = np.zeros(1, dtype=np.float64)

E2_PE = np.zeros(1, dtype=np.float64)
E2_PX = np.zeros(1, dtype=np.float64)
E2_PY = np.zeros(1, dtype=np.float64)
E2_PZ = np.zeros(1, dtype=np.float64)

Jpsi_M = np.zeros(1, dtype=np.float64)

my_tree.SetBranchAddress('E1_PE', E1_PE)
my_tree.SetBranchAddress('E1_PX', E1_PX)
my_tree.SetBranchAddress('E1_PY', E1_PY)
my_tree.SetBranchAddress('E1_PZ', E1_PZ)

my_tree.SetBranchAddress('E2_PE', E2_PE)
my_tree.SetBranchAddress('E2_PX', E2_PX)
my_tree.SetBranchAddress('E2_PY', E2_PY)
my_tree.SetBranchAddress('E2_PZ', E2_PZ)

newbranch1 = new_tree.Branch('Jpsi_M', Jpsi_M, 'Jpsi_M/D')

for i in range(my_tree.GetEntries()):
    my_tree.GetEntry(i)
    PE = E1_PE[0] + E2_PE[0]
    PX = E1_PX[0] + E2_PX[0]
    PY = E1_PY[0] + E2_PY[0]
    PZ = E1_PZ[0] + E2_PZ[0]
    Jpsi_M[0] = np.sqrt(PE**2 - PX**2 - PY**2 - PZ**2)
    newbranch1.Fill()

### Aufgabenteil d)
N_before = new_tree.GetEntries()
N_after = new_tree.GetEntries('Jpsi_M>2400&&Jpsi_M<3300')
eff = N_after/N_before
err = np.sqrt(eff*(1 - eff)/N_before)
print("The efficiency is: {:.3f}+/-{:.3f}".format(eff, err))

new_tree.Write('', R.TObject.kOverwrite)
new_file.Close()

my_file.Close()