#!/usr/bin/python
import os,sys
from ROOT import *

from robs_draw_fiducial import *
ls = fiducial_lines()

fid = [0.54, 0.96,  52.0, 10.0,   14.0,      38.0] ##intermediate fiducial

fn = 'sum74-4-1_sum474-1_2_21-1_michel.root'
f = TFile(fn)

#from /home/alexg/residual_ana/fidu-test/binning00.py
## this is 6D already
def in_fid(p,ct,cvals):
    if not cvals[0]<abs(ct)<cvals[1]: return False
    if not p<cvals[2]: return False
    pz=p*ct
    if not abs(pz)>cvals[4]: return False
    pt=sqrt(p**2-pz**2)
    if not cvals[3]<pt<cvals[5]: return False
    return True

###########################################################################

h_all_residuals = f.Get('full_residuals2').Clone()
# h_resid_ptot    = f.Get('full_residuals_ptot').Clone()
# h_resid_costh   = f.Get('full_residuals_costh').Clone()

c1 = TCanvas('c1','c1',850,1100) #letter ratio
c1.Divide(1,2)

c1.cd(1)
c1_1.Divide(2,1)
c1_1.cd(1)

h_all_residuals.SetTitle('Normalised residuals; momentum (MeV/c); cos#theta')
h_all_residuals.GetYaxis().SetDecimals()
h_all_residuals.Clone().Draw('colz')
for l in ls: l.Draw('same')

#c1.cd(2)
c1_1.cd(2)
h_all_residuals.Rebin2D()
h_all_residuals.Rebin2D()
h_all_residuals.SetMaximum(3.)
h_all_residuals.SetMinimum(-3.)
h_all_residuals.Clone().Draw('colz')
#for l in ls: l.Draw('same')

def list_residuals(h,ctmin,ctmax):
#DS = 0<ct<1, US = -1<ct<0
    rs = []
    nx,ny = h.GetNbinsX(),h.GetNbinsY()
    for ibx in range(1,nx+1):
        p  = h.GetXaxis().GetBinCenter(ibx)
        for iby in range(1,ny+1):
            ct = h.GetYaxis().GetBinCenter(iby)
            if in_fid(p,ct,fid) and ctmin<ct<ctmax:
                rs.append(h.GetBinContent(ibx,iby))
    return rs

#h_residuals_fid_only = fid_only(h_all_residuals,fid,999.,999.)
#h_residuals_fid_only.Draw('colz')

US = list_residuals(h_all_residuals,0.,1.)
DS = list_residuals(h_all_residuals,-1.,0.)
US_and_DS = US+DS

h_DS        = TH1D('h_DS','h_DS',50,-5.,5.)
h_US        = TH1D('h_US','h_US',50,-5.,5.)
h_US_and_DS = TH1D('h_US_and_DS','h_US_and_DS',50,-5.,5.)

for r in US:
    h_US.Fill(r,1.)
    h_US_and_DS.Fill(r,1.)
for r in DS:
    h_DS.Fill(r,1.)
    h_US_and_DS.Fill(r,1.)

h_US_and_DS.Fit('gaus')
h_US_and_DS.SetTitle('Residuals in fiducial only; (data-simulation)/#sigma;')
h_US_and_DS.Draw()


#####################################################################################
#now the residual projections

# #first bit is crap to deal with root splitting canvii
# fudge = 2.0
# a = [h_resid_ptot.GetXaxis(),h_resid_ptot.GetYaxis(),
#      h_resid_costh.GetXaxis(),h_resid_costh.GetYaxis()]
# for ai in a:
#     ai.SetTitleSize(ai.GetTitleSize()*fudge)
#     ai.SetLabelSize(ai.GetLabelSize()*fudge)
# 
# h_resid_ptot.GetYaxis().SetNdivisions(505)
# h_resid_costh.GetYaxis().SetNdivisions(505)
# 
# pmin_cut = TLine(17.3,-5.,17.3,5.)
# pmax_cut = TLine(52.0,-5.,52.0,5.)
# ctmin_cut1 = TLine(0.54,-5.,0.54,5.)
# ctmin_cut2 = TLine(-0.54,-5.,-0.54,5.)
# ctmax_cut1 = TLine(0.96,-5.,0.96,5.)
# ctmax_cut2 = TLine(-0.96,-5.,-0.96,5.)
# for cs in [pmin_cut,pmax_cut,ctmin_cut1,ctmax_cut1,ctmin_cut2,ctmax_cut2]:
#     cs.SetLineColor(4)#blue
#     cs.SetLineWidth(5) #to match top left plot
# 
# c1.cd(2)
# c1_2.Divide(1,2)
# c1_2.cd(1)
# #c1_3.Divide(2,1)
# #c1_3.cd(1)
# #c1_3.SetGridx()
# #c1_3.SetGridy()
# h_resid_ptot.SetMinimum(-5.)
# h_resid_ptot.SetMaximum(5.)
# h_resid_ptot.SetTitle('; momentum (MeV/c); (data-simulation)/#sigma')
# h_resid_ptot.GetXaxis().SetRangeUser(15,60)
# 
# h_resid_ptot.Draw()
# pmin_cut.Draw('same')
# pmax_cut.Draw('same')
# 
# c1_2.cd(2)
# #c1_4.SetGridx()
# #c1_4.SetGridy()
# 
# #c1_4.Divide(2,1)
# #c1_4.cd(1)
# 
# h_resid_costh.SetMinimum(-5.)
# h_resid_costh.SetMaximum(5.)
# h_resid_costh.SetTitle('; cos#theta; (data-simulation)/#sigma')
# h_resid_costh.GetXaxis().SetDecimals()
# #h_resid_costh.GetXaxis().SetRangeUser(-1.0,-0.3)
# h_resid_costh.Clone().Draw()
# 
# h_resid_costh.GetYaxis().SetNdivisions(505)
# 
# ctmin_cut2.Draw('same')
# ctmax_cut2.Draw('same')
# 
# #c1_4.cd(2)
# #h_resid_costh.GetXaxis().SetRangeUser(0.3,1.0)
# #h_resid_costh.Clone().Draw()
# 
# ctmin_cut1.Draw('same')
# ctmax_cut1.Draw('same')
# 
au = raw_input('>')