#!/usr/bin/python
import os,sys
from ROOT import *
from clean_canvas import *
from nice_colours import *
gStyle = makegstyle(gStyle)
gStyle = nice_colours(gStyle)
gStyle.SetOptFit(1011)
gStyle.SetErrorX(0)
gStyle.SetStatX(0.88)
gStyle.SetStatY(0.85)
gStyle.SetStatW(0.4)
gStyle.SetStatH(0.15)

from robs_draw_fiducial import *
ls = fiducial_lines()

from accessories import *
fid = [0.54, 0.96,  52.0, 10.0,   14.0,      38.0] ##intermediate fiducial

#fn = '/home/e614/analysis_2006_2007/treesum/pass4/sum74-4-1_sum474-1_2_21-1_michel.root'
fn = '/home/e614/analysis_2006_2007/treesum/pass8/sum87-3-1_sum587-1_21-1_michel.root'  #used for plot
#fn = '/home/e614/analysis_2006_2007/treesum/systematics_pass1/DistortedSTRs_ECal_AscaleBscale/sum584-3_4-1_sum584-1_2-1_michel.root' #distorted STRS
#fn = '/home/e614/analysis_2006_2007/treesum/systematics_pass1/sum440-1-1_sum474-1_2_21_31-1_michel.root' #brem
#fn = '/home/e614/analysis_2006_2007/treesum/systematics_pass1/sum445-1-1_sum474-1_2_21_31-1_michel.root' #foil pos
#fn = '/home/e614/analysis_2006_2007/treesum/systematics_pass1/EtaCorrelation/sum84-9-4_sum584-1_2_21-1_etaup_michel.root' #eta up

f = TFile(fn)
nice_fn = fn.split('/')[-1].rstrip('.root')

def zero_y_errors(h1):
    h2 = h1.Clone()
    h2.Reset()
    for ib in range(1,h1.GetNbinsX()+1):
        c = h1.GetBinCenter(ib)
        n = h1.GetBinContent(ib)
        e = h1.GetBinError(ib)
        if n!=0.0 and e!=0.0:
            h2.SetBinContent(ib,n)
            h2.SetBinError(ib,0.)
        else:
            h2.SetBinContent(ib,-1000.) #so embarrassing
            h2.SetBinError(ib,0.)
    return h2

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

for h in [h_resid_ptot,h_resid_costh]:
    h.SetMarkerStyle(20)
    h.SetMarkerColor(1)
    h.SetMarkerSize(1.)

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

c1.cd(1)
c1_1.Divide(2,1)
c1_1.GetPad(1).SetLeftMargin(.12)
c1_1.GetPad(1).SetRightMargin(.15)
c1_1.GetPad(2).SetLogy()
c1_1.cd(1)

h_all_residuals.SetTitle('Normalised residuals; momentum (MeV/c); cos#theta')
h_all_residuals.GetYaxis().SetDecimals()
h_all_residuals.GetYaxis().CenterTitle()
h_all_residuals.GetYaxis().SetTitleOffset(1.6)
h_all_residuals.SetMinimum(-10.)
h_all_residuals.SetMaximum(10.)
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.SetMaximum(1E4)
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)
for p in [1,2]:
    c1_2.GetPad(p).SetGridx()
    c1_2.GetPad(p).SetGridy()
    c1_2.GetPad(p).SetBottomMargin(.2)
    c1_2.GetPad(p).SetLeftMargin(.07)

c1_2.cd(1)

h_resid_ptot.SetMinimum(-5.)
h_resid_ptot.SetMaximum(5.)
#h_resid_ptot.SetTitle(nice_fn+'; momentum (MeV/c); (data-simulation)/#sigma')
h_resid_ptot.SetTitle('; momentum (MeV/c); (data-simulation)/#sigma')
h_resid_ptot.GetXaxis().SetRangeUser(15,60)
h_resid_ptot.GetYaxis().SetTitleOffset(0.4)

h_resid_ptot.Draw('p')
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.GetYaxis().SetTitleOffset(0.4)
#h_resid_costh.GetXaxis().SetRangeUser(-1.0,-0.3)
h_resid_costh.Clone().Draw('p')

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')

c1.Update()
au = raw_input('>')