#!/usr/bin/env python
import sys
from ROOT import gROOT, gStyle, TFile, TTree, TCanvas, gDirectory, TH1, TGraph, gPad, TF1

c = TCanvas("c","c",1200,900);
c.Print("low_x.pdf[")

gROOT.SetBatch(True)
gStyle.SetOptFit(1)
gStyle.SetStatW(0.2)
gStyle.SetStatH(0.1)

# The "magic cut" that cleans up the mass distribution is using truth information to require p(A')>0.8*Ebeam (in addition to the same cut on the reconstructed pair). This gives you a much more symmetric mass distribution.
# The reason this cut works is that there is a significant population of A' generated at low x, where the detector and recon pick up the daughter positron and the recoil e-.
# This contamination is O(10%) in these plots (with minimal cuts), probably less after standard analysis cuts.

# This magic cut is obviously cheating, since there's no simple/clean equivalent cut in data.
# In my thesis analysis I used this cut when calculating acceptance and mass resolution.
# The main reason was that without the cut, the acceptance vs. Z fits are harder (the decay Z distribution is no longer a clean exponential because the A' don't all have the same boost)
# The right thing to do is to replace this with a mass cut.

# It might be possible to find a cut in data specific to these events (maybe in vertexing? the recoil e- will point back to the target, the e+ will not - similar to a L1 WAB conversion) but not easy.

# this file is available at http://www.slac.stanford.edu/~meeg/hps3/users/meeg/vertexing/acceptance/apsignal_displaced_60_dq_tri.root
# 60 MeV A', truth tuple after reconstruction
inFile = TFile("apsignal_displaced_60_dq_tri.root")
ntuple = inFile.Get("ntuple")

ntuple.SetAlias("corrM","uncM-0.15e-3*(elePX/eleP-posPX/posP)*uncVZ/uncM")
#"-0.15e-3*(elePX/eleP-posPX/posP)*uncVZ/uncM" is an empirical correction to the misreconstructed mass from the vertexer
#triEndZ is the true decay Z, a cut at 5 selects near-prompt A'
#triP is the truth momentum of the 622 particle (i.e. the A' momentum thrown by MadGraph)

ntuple.Draw("corrM-0.06:triEndZ>>h(100,-10,100,100,-0.04,0.04)","uncP>0.8*1.056","colz")
gDirectory.Get("h").SetTitle("uncP/Ebeam>0.8;decay Z [mm];mass residual [GeV]")
gPad.SetLogy(0)
c.Print("low_x.pdf")

ntuple.Draw("corrM-0.06:triEndZ>>h(100,-10,100,100,-0.04,0.04)","triP>0.8*1.056&&uncP>0.8*1.056","colz")
gDirectory.Get("h").SetTitle("p(A')/Ebeam>0.8, uncP/Ebeam>0.8;decay Z [mm];mass residual [GeV]")
gPad.SetLogy(0)
c.Print("low_x.pdf")

# anticut
ntuple.Draw("corrM-0.06:triEndZ>>h(100,-10,100,100,-0.04,0.04)","triP<0.8*1.056&&uncP>0.8*1.056","colz")
gDirectory.Get("h").SetTitle("p(A')/Ebeam<0.8, uncP/Ebeam>0.8;decay Z [mm];mass residual [GeV]")
gPad.SetLogy(0)
c.Print("low_x.pdf")

ntuple.Draw("corrM-0.06>>h(100,-0.04,0.04)","triEndZ<5&&uncP>0.8*1.056","colz")
gDirectory.Get("h").SetTitle("prompt A', uncP/Ebeam>0.8;mass residual [GeV]")
gDirectory.Get("h").Fit("gaus","","",-0.005,0.005)
gPad.SetLogy(1)
c.Print("low_x.pdf")

# after the magic cut
ntuple.Draw("corrM-0.06>>h(100,-0.04,0.04)","triEndZ<5&&triP>0.8*1.056&&uncP>0.8*1.056","colz")
gDirectory.Get("h").SetTitle("prompt A', p(A')/Ebeam>0.8, uncP/Ebeam>0.8;mass residual [GeV]")
gDirectory.Get("h").Fit("gaus","","",-0.005,0.005)
gPad.SetLogy(1)
c.Print("low_x.pdf")

# anticut
ntuple.Draw("corrM-0.06>>h(100,-0.04,0.04)","triEndZ<5&&triP<0.8*1.056&&uncP>0.8*1.056","colz")
gDirectory.Get("h").SetTitle("prompt A', p(A')/Ebeam<0.8, uncP/Ebeam>0.8;mass residual [GeV]")
gDirectory.Get("h").Fit("gaus","","",-0.005,0.005)
gPad.SetLogy(1)
c.Print("low_x.pdf")


gStyle.SetOptStat(0)

ntuple.Draw("corrM-0.06:triP/1.056>>h(100,0,1.1,100,-0.04,0.04)","triEndZ<5&&uncP>0.8*1.056","colz")
gDirectory.Get("h").SetTitle("prompt A', uncP/Ebeam>0.8;p(A')/Ebeam;mass residual [GeV]")
gPad.SetLogy(0)
c.Print("low_x.pdf")


# high p(A') events have good correlation between truth and recon p(e+e-); low p(A') events are uncorrelated
ntuple.Draw("uncP/1.056:triP/1.056>>h(100,0,1.1,100,0.5,1.1)","triEndZ<5&&uncP>0.8*1.056","colz")
gDirectory.Get("h").SetTitle("prompt A', uncP/Ebeam>0.8;p(A')/Ebeam;uncP/Ebeam")
gPad.SetLogy(0)
c.Print("low_x.pdf")

ntuple.Draw("uncP/1.056:triP/1.056>>h(100,0,1.1,100,0,1.1)","triEndZ<5","colz")
gDirectory.Get("h").SetTitle("prompt A', uncP/Ebeam>0.8;p(A')/Ebeam;uncP/Ebeam")
gPad.SetLogy(0)
c.Print("low_x.pdf")

ntuple.Draw("triPosP/1.056:triP/1.056>>h(100,0,1.1,100,0,1.1)","triEndZ<5&&uncP>0.8*1.056","colz")
gDirectory.Get("h").SetTitle("prompt A', uncP/Ebeam>0.8;p(A')/Ebeam;p(e+ from A' decay)/Ebeam")
gPad.SetLogy(0)
c.Print("low_x.pdf")

ntuple.Draw("triEle1P/1.056:triP/1.056>>h(100,0,1.1,100,0,1.1)","triEndZ<5&&uncP>0.8*1.056","colz")
gDirectory.Get("h").SetTitle("prompt A', uncP/Ebeam>0.8;p(A')/Ebeam;p(e- from A' decay)/Ebeam")
gPad.SetLogy(0)
c.Print("low_x.pdf")

# in low p(A') events, p(recon e-)+p(a')=Ebeam, this proves we're picking up the recoil e-
ntuple.Draw("eleP/1.056:triP/1.056>>h(100,0,1.1,100,0,1.1)","triEndZ<5&&uncP>0.8*1.056","colz")
gDirectory.Get("h").SetTitle("prompt A', uncP/Ebeam>0.8;p(A')/Ebeam;p(reconstructed e-)/Ebeam")
gPad.SetLogy(0)
c.Print("low_x.pdf")

# two populations of events: truth and recon p(e-) match, or truth p(e-) is much lower (we are picking the wrong e-)
ntuple.Draw("triEle1P/1.056:eleP/1.056>>h(100,0,1.1,100,0,1.1)","triEndZ<5&&uncP>0.8*1.056","colz")
gDirectory.Get("h").SetTitle("prompt A', uncP/Ebeam>0.8;p(reconstructed e-)/Ebeam;p(e- from A' decay)/Ebeam")
gPad.SetLogy(0)
c.Print("low_x.pdf")

# use the difference between truth and recon p(e-) as a proxy for "are we picking the correct e-": events with the wrong e- have broad, negative mass residual distribution
ntuple.Draw("corrM-0.06:(eleP-triEle1P)/1.056>>h(100,-0.5,1.1,100,-0.04,0.04)","triEndZ<5&&uncP>0.8*1.056","colz")
gDirectory.Get("h").SetTitle("prompt A', uncP/Ebeam>0.8;p(reconstructed e-)-p(e- from A' decay)/Ebeam;mass residual [GeV]")
gPad.SetLogy(0)
c.Print("low_x.pdf")


c.Print("low_x.pdf]")