#!/usr/bin/python
import os,sys
from ROOT import *
from clean_canvas import *
from math import sqrt
gStyle = makegstyle(gStyle)
gStyle.SetOptFit(1)
#rbin = 1
rbin = 8

wd = '/triumfcs/trshare/alexg/pos-int/'       #working directory
fns = {'data':'s74delta105pc.root',
       'sim':'g474delta105pc.root',
       'delta_boost_3':'g440delta105pc.root',
       'brem_boost_3':'g441delta105pc.root',
       'delta_boost_0':'g442delta105pc.root'}
fs = []

def retrieve_histo(fn,hn,ti):
    if not os.path.isfile(wd+fn):
        print 'ERROR: cannot find',wd+fn
        sys.exit(-1)
    fs.append(TFile(wd+fn))
    h = fs[-1].Get(hn).Clone()
    h.SetTitle(ti)
    return h

def renorm_factor(fn1,fn2):
    if not os.path.isfile(wd+fn1) or not os.path.isfile(wd+fn2):
        print 'ERROR: not a file:',fn1,fn2
        sys.exit(-1)
    f1 = TFile(wd+fn1)
    h1 = f1.Get('TopologyCut/ntracks')
    n1 = h1.GetEntries()
    f2 = TFile(wd+fn2)
    h2 = f2.Get('TopologyCut/ntracks')
    n2 = h2.GetEntries()
    return n1/n2

def take_ratio_histograms(h1,h2):
    nb1 = h1.GetNbinsX()
    nb2 = h2.GetNbinsX()
    if nb1!=nb2:
        print 'ERROR: histograms do not have same number of x bins'
        sys.exit(-1)
    h3 = h2.Clone()
    h3.Reset()
    h3.SetTitle('ratio of '+h1.GetTitle()+' and '+h2.GetTitle())
    h3.GetYaxis().SetTitle('ratio')
    for ix in range(nb1):
        n1 = h1.GetBinContent(ix)
        n2 = h2.GetBinContent(ix)
        if n1>0. and n2>0.:
#            print ix,n1,n2,n1/n2
            h3.SetBinContent(ix,(n1/n2))
#            h3.SetBinError(ix,(n1/n2)*sqrt( (sqrt(n1)/n1)**2 + (sqrt(n2)/n2)**2 ))
            h3.SetBinError(ix,(n1/n2)*sqrt( 1/n1 + 1/n2 ))  #equivalent to line above
        else:
            h3.SetBinContent(ix,0.)
            h3.SetBinError(ix,0.)
    return h3.Clone()
#

alabels = "; #delta momentum (MeV/c); reconstructed #delta's"

def count_entries(h,pmin,pmax):
    t = 0.
    nb = h.GetNbinsX()
    for i in range(nb):
        if pmin < h.GetBinCenter(i) < pmax:
            t += h.GetBinContent(i)
    return t
def count_entries_US(h,pmin,pmax):
    #this still doesn't really work
    t = 0.
    nb = h.GetNbinsX()
    for i in range(nb):
        pcent =  h.GetBinCenter(i)
        if (pmin < pcent < 18.0) or (23.5 < pcent < pmax):
            t += h.GetBinContent(i)
    return t

def style_histo(h):
    for a in [h.GetXaxis(),h.GetYaxis()]:
        a.SetTitleOffset(1.2)
        a.SetTitleSize(0.06)
#        a.SetTitleFont(42)
        a.SetLabelSize(0.06)
#        a.SetLabelFont(42)
#    h.SetTitleFont(42)
    return h

def make_legend(h1,hti1,h2,hti2):
    l = TLegend(.8,.8,.9,.9)
    l.SetBorderSize(1)
    l.AddEntry(h1,hti1,'l')
    l.AddEntry(h2,hti2,'l')
    return l.Clone()

def do_damn_analysis(fn1,fn2,atype):
    if atype == 'target':
        dplot = 'TopologyCut/pp2ds'
#        dplot = 'TopologyCut/pp2us'
        pmin_fit = 6
        pmax_fit = 26
#        pmax_fit = 36
    elif atype == 'DC':
        dplot = 'TopologyCut/pp2diff3'
#        dplot = 'TopologyCut/pp2diff6'
        pmin_fit = 6
        pmax_fit = 26
#        pmax_fit = 36
    elif atype == 'brem':
        dplot = 'TopologyCut/pp2diff'
        pmin_fit = 6
        pmax_fit = 36
    else:
        print 'ERROR: analysis must be target or DC'
        sys.exit(-1)
    h1 = retrieve_histo(fn1,dplot,'Simulation: '+atype+' deltas, DS'+alabels)
    h2 = retrieve_histo(fn2,dplot,'Data: '+atype+' deltas, DS'+alabels)
    h1 = style_histo(h1)
    h2 = style_histo(h2)
    h2.SetLineStyle(7)
    h2.Scale(renorm_factor(fn1,fn2))
    n1 = count_entries(h1,pmin_fit,pmax_fit)
    n2 = count_entries(h2,pmin_fit,pmax_fit)
#    n1 = count_entries_US(h1,pmin_fit,pmax_fit)
#    n2 = count_entries_US(h2,pmin_fit,pmax_fit)
    print
    print fn1,':',n1
    print fn2,':',n2
    print 'ratio=',n2/n1,'+-',(n2/n1)*sqrt(1/n1+1/n2)
    print
    h1.Rebin(rbin)
    h2.Rebin(rbin)
    r12 = take_ratio_histograms(h2,h1)
    h1.GetXaxis().SetRangeUser(-10.,60.)
    h2.GetXaxis().SetRangeUser(-10.,60.)
    r12.GetXaxis().SetRangeUser(-10.,60.)
    return h1,h2,r12

print '__________ delta target __________'

c_target = TCanvas('c_target','c_target',0,0,1000,1150)
c_target.Divide(2,3)
for ci in [c_target_1,c_target_2,c_target_3,c_target_4,c_target_5,c_target_6]:
    ci.SetLeftMargin(.15)
    ci.SetRightMargin(.02)
    ci.SetBottomMargin(.13)
    ci.SetGridx()
    ci.SetGridy()

#ds = data_sim
ds_1, ds_2, ds_r12 = do_damn_analysis(fns['sim'],fns['data'],'target')
c_target.cd(1)
ds_1.Draw()
ds_2.Draw('same')
l = make_legend(ds_1,'simulation (g474)',ds_2,'data (s74)')
l.Draw()
c_target.cd(2)
ds_r12.Draw()

#s3 = sim_sim3
s3_1, s3_2, s3_r12 = do_damn_analysis(fns['sim'],fns['delta_boost_3'],'target')
c_target.cd(3)
s3_1.Draw()
s3_2.Draw('same')
l = make_legend(s3_1,'simulation (g474)',s3_2,'#delta rate x 3')
l.Draw()
c_target.cd(4)
s3_r12.Draw()

#s0 = sim_sim0
s0_1, s0_2, s0_r12 = do_damn_analysis(fns['sim'],fns['delta_boost_0'],'target')
c_target.cd(5)
s0_1.Draw()
s0_2.Draw('same')
l = make_legend(s0_1,'simulation (g474)',s0_2,"#delta's off")
l.Draw()
c_target.cd(6)
s0_r12.Draw()

print '__________ delta DCs __________'

c_DC = TCanvas('c_DC','c_DC',0,0,1000,1150)
c_DC.Divide(2,3)
for ci in [c_DC_1,c_DC_2,c_DC_3,c_DC_4,c_DC_5,c_DC_6]:
    ci.SetLeftMargin(.15)
    ci.SetRightMargin(.02)
    ci.SetBottomMargin(.13)
    ci.SetGridx()
    ci.SetGridy()

#ds = data_sim
dc_ds_1, dc_ds_2, dc_ds_r12 = do_damn_analysis(fns['sim'],fns['data'],'DC')
dc_ds_2.SetLineStyle(7)
c_DC.cd(1)
dc_ds_1.Draw()
dc_ds_2.Draw('same')
l = make_legend(dc_ds_1,'simulation (g474)',dc_ds_2,'data (s74)')
l.Draw()
c_DC.cd(2)
dc_ds_r12.Draw()

#s3 = sim_sim3
dc_s3_1, dc_s3_2, dc_s3_r12 = do_damn_analysis(fns['sim'],fns['delta_boost_3'],'DC')
c_DC.cd(3)
dc_s3_1.Draw()
dc_s3_2.Draw('same')
l = make_legend(dc_s3_1,'simulation (g474)',dc_s3_2,'#delta rate x 3')
l.Draw()
c_DC.cd(4)
dc_s3_r12.Draw()

#s0 = sim_sim0
dc_s0_1, dc_s0_2, dc_s0_r12 = do_damn_analysis(fns['sim'],fns['delta_boost_0'],'DC')
c_DC.cd(5)
dc_s0_1.Draw()
dc_s0_2.Draw('same')
l = make_legend(dc_s0_1,'simulation (g474)',dc_s0_2,"#delta's off")
l.Draw()
c_DC.cd(6)
dc_s0_r12.Draw()

print '__________ brem __________'

c_brem = TCanvas('c_brem','c_brem',0,0,1000,1150)
c_brem.Divide(2,3)
for ci in [c_brem_1,c_brem_2,c_brem_3,c_brem_4,c_brem_5,c_brem_6]:
    ci.SetLeftMargin(.15)
    ci.SetRightMargin(.02)
    ci.SetBottomMargin(.13)
    ci.SetGridx()
    ci.SetGridy()

#ds = data_sim
brem_ds_1, brem_ds_2, brem_ds_r12 = do_damn_analysis(fns['sim'],fns['data'],'brem')
c_brem.cd(1)
brem_ds_1.Draw()
brem_ds_2.Draw('same')
c_brem.cd(2)
brem_ds_r12.Draw()

#s3 = sim_sim3
brem_s3_1, brem_s3_2, brem_s3_r12 = do_damn_analysis(fns['sim'],fns['brem_boost_3'],'brem')
c_brem.cd(3)
brem_s3_1.Draw()
brem_s3_2.Draw('same')
c_brem.cd(4)
brem_s3_r12.Draw()

print '__________ tail count delta __________'

c_delta_tail = TCanvas('c_delta_tail','c_delta_tail',0,0,1000,1150)
c_delta_tail.Divide(2,3)
for ci in [c_delta_tail_1,c_delta_tail_2,c_delta_tail_3,c_delta_tail_4,c_delta_tail_5,c_delta_tail_6]:
    ci.SetLeftMargin(.15)
    ci.SetRightMargin(.02)
    ci.SetBottomMargin(.13)
    ci.SetGridx()
    ci.SetGridy()

#ds = data_sim
#tail_1, tail_2, tail_r12 = do_damn_analysis(fns['delta_boost_0'],fns['delta_boost_3'],'brem')
tail_1, tail_2, tail_r12 = do_damn_analysis(fns['sim'],fns['delta_boost_3'],'brem')
c_delta_tail.cd(1)
tail_1.Draw()
tail_2.Draw('same')
c_delta_tail.cd(2)
tail_r12.Draw()

au = raw_input('>')