#!/usr/bin/python

from ROOT import TFile, TH1, TGraph, TCanvas, TF1, TPad, TGraphErrors
from ROOT import kFullCircle, kFullSquare, kBlack, kBlue, kOpenSquare, kOpenCircle
from ROOT import gPad
import numpy as np
from array import array
from physics_calcs import *

def DrawBackScatteringParams() :
    c = r.TCanvas("c1", "c1", 800, 800)
    c.Divide(1, 3)

    c.cd(1)
    x1 = array('d', [1, 10, 100, 1000])
    sim1 = array('d', [1.35, 1.32, 1.25, 0.81])
    dsim1 = array('d', [.006, .01, .01, .01])    
    zero = array('d', [0]*len(x1))

    gdata1 = AddHatchedRegion(0, 1000, 1.39, 1.43)
    gsim1 = TGraphErrors(len(x1), x1, sim1, zero, dsim1)
    gsim1.GetXaxis().SetTitle("Cut For Secondaries [#mum]")
    gsim1.GetYaxis().SetTitle("Backscattering fraction")
    gsim1.Draw("AP")
    gPad.Update()
    gdata1.Draw("F SAME")
    gPad.SetLogx(1)

    FormatAxis(gsim1, ytsize = .09, ytoffset = 0.5, ytfont = 132,
               ylsize = 0.08, ylfont = 132, xtsize = 0.1, xtoffset = 0.9,
               xlsize = 0.09, xloffset = -0.013)
    gsim1.GetXaxis().SetLabelOffset(-0.015)
    gPad.Update()
    gsim1.GetYaxis().SetRangeUser(0.65, 1.6)
    gPad.Update()
    c.cd(2)
    x2 = array('d', [.002, .02, .04])
    sim2 = array('d', [1.35, 1.30, 1.32])
    dsim2 = array('d', [.006, .01, .01])
    zero2 = array('d', [0]*len(x2))
    gsim2 = TGraphErrors(len(x2), x2, sim2, zero2, dsim2)
    gsim2.GetXaxis().SetTitle("Range Factor")
    gsim2.GetYaxis().SetTitle("Backscattering Fraction")
    FormatAxis(gsim2, ytsize = .09, ytoffset = 0.5, ytfont = 132,
               ylsize = 0.08, ylfont = 132, xtsize = 0.1, xtoffset = 0.9,
               xlsize = 0.09, xloffset = -0.013)
    gsim2.Draw("AP")
    gdata2 = AddHatchedRegion(0, .05, 1.39, 1.43)
    gdata2.Draw("F SAME")
    gPad.Update()
    gsim2.GetYaxis().SetRangeUser(.65, 1.6)
    gPad.Update()
    c.cd(3)
    x3 = array('d', [1, 2, 3])
    sim3 = array('d', [1.35, 1.34, 1.32])
    dsim3 = array('d', [.006, .01, .01])
    gsim3 = TGraphErrors(len(x3), x3, sim3, zero2, dsim3)
    gsim3.GetYaxis().SetTitle("Backscattering Fraction")
    FormatAxis(gsim3, ytsize = .09, ytoffset = 0.5, ytfont = 132,
               ylsize = 0.08, ylfont = 132, xtsize = 0., xlsize = 0.)
    gsim3.Draw("AP")
    gsim3.GetXaxis().SetTickLength(0)
    gPad.Update()
    gdata3 = AddHatchedRegion(0, 4, 1.39, 1.43)
    gdata3.Draw("F SAME")
    gsim3.GetYaxis().SetRangeUser(.65, 1.6)
    gPad.Update()
    raw_input("enter to end")
def AddHatchedRegion(xmin, xmax, ymin, ymax) :
    n = 20
    x = array('d', [0]*n)
    ymin = array('d', [ymin]*n)
    ymax = array('d', [ymax]*n)
    for i in range(n) :
        x[i] = xmin + i*((xmax-xmin)/float(n-1))
    grmin = TGraph(n, x, ymin)
    grmax = TGraph(n, x, ymax)
    grshade = TGraph(2*n)
    for i in range(n) :
        grshade.SetPoint(i, x[i], ymax[i])
        grshade.SetPoint(n+i,x[n-i-1],ymin[n-i-1])

    grshade.SetFillStyle(3013)
    grshade.SetFillColor(kBlue)
    return grshade

def SimpleScintSpectrumWithBB1(tree) :
    cutu = "(BB1_AMPLITUDE_UX[0] > 16"
    cutl = "(BB1_AMPLITUDE_LX[0] > 16"
    for i in range(1, 40) :
        cutu = "%s || BB1_AMPLITUDE_UX[%d] > 16" % (cutu, i)
        cutl = "%s || BB1_AMPLITUDE_LX[%d] > 16" % (cutl, i)
    cutu = cutu + ") && (BB1_AMPLITUDE_UY[0] > 16"
    cutl = cutl + ") && (BB1_AMPLITUDE_LY[0] > 16"
    for i in range(1, 40) :
        cutu = "%s || BB1_AMPLITUDE_UY[%d] > 16" % (cutu, i)
        cutl = "%s || BB1_AMPLITUDE_LY[%d] > 16" % (cutl, i)
    cutu = cutu + ")"
    cutl = cutl + ")"

    print cutu
    print cutl

    tree.Draw("QDC_UpperPMT>>hup(4095, -0.5, 4095.5)", cutu)
    tree.Draw("QDC_LowerPMT>>hlo(4095, -0.5, 4095.5)", cutl, "SAME")
    gPad.Update()
    raw_input("enter to end")


def CompareBackscattering(fname1, title1, fname2, title2) :
    f1 = TFile(fname1)
    f2 = TFile(fname2)

    su1 = f1.Get("Upper_single_scint_E")
    du1 = f1.Get("Upper_double_scint_E")
    
    sl1 = f1.Get("Lower_single_scint_E")
    dl1 = f1.Get("Lower_double_scint_E")

    su2 = f2.Get("Upper_single_scint_E")
    du2 = f2.Get("Upper_double_scint_E")
    
    sl2 = f2.Get("Lower_single_scint_E")
    dl2 = f2.Get("Lower_double_scint_E")

    # Get fraction for 1000 keV to 5100 keV
    emin = [1000., 1500.]
    print "Integrated BS fractions for ", fname1
    print "Upper detector:"
    for e in emin :
        single1_up = su1.Integral(su1.FindBin(e), su1.FindBin(5100.))
        double1_up = du1.Integral(du1.FindBin(e), du1.FindBin(5100.))
        perc1_up = 100.*(ufloat(double1_up, sqrt(double1_up))/ufloat(single1_up, sqrt(single1_up)))

        print "\tPercent ", e, " keV", perc1_up
    print "Lower detector:"
    for e in emin :
        single1_up = sl1.Integral(sl1.FindBin(e), sl1.FindBin(5100.))
        double1_up = dl1.Integral(dl1.FindBin(e), dl1.FindBin(5100.))
        perc1_up = 100.*(ufloat(double1_up, sqrt(double1_up))/ufloat(single1_up, sqrt(single1_up)))

        print "\tPercent ", e, " keV", perc1_up
    print "Integrated BS fractions for ", fname2
    print "Upper detector:"
    for e in emin :
        single1_up = su2.Integral(su2.FindBin(e), su2.FindBin(5100.))
        double1_up = du2.Integral(du2.FindBin(e), du2.FindBin(5100.))
        perc1_up = 100.*(ufloat(double1_up, sqrt(double1_up))/ufloat(single1_up, sqrt(single1_up)))

        print "\tPercent ", e, " keV", perc1_up
    print "Lower detector:"
    for e in emin :
        single1_up = sl2.Integral(sl2.FindBin(e), sl2.FindBin(5100.))
        double1_up = dl2.Integral(dl2.FindBin(e), dl2.FindBin(5100.))
        perc1_up = 100.*(ufloat(double1_up, sqrt(double1_up))/ufloat(single1_up, sqrt(single1_up)))

        print "\tPercent ", e, " keV", perc1_up

    rbin = 256
    for h in [su1, du1, sl1, dl1, su2, du2, sl2, dl2] :
        h.Rebin(rbin)
        h.Sumw2()

    upper_1 = du1.Clone()
    lower_1 = dl1.Clone()
    upper_2 = du2.Clone()
    lower_2 = dl2.Clone()

    upper_1.Divide(su1)
    lower_1.Divide(sl1)
    upper_2.Divide(su2)
    lower_2.Divide(sl2)


    upper_1.SetTitle(title1 + " Upper")
    lower_1.SetTitle(title1 + " Lower")
    upper_2.SetTitle(title2 + " Upper")
    lower_2.SetTitle(title2 + " Lower")

    c1 = TCanvas("c1", "c1", 1200, 600)
    c1.SetBottomMargin(10*c1.GetBottomMargin())
    p1 = TPad("p1", "p1", 0., 0.35 , 1., 1.  )
    p2 = TPad("p2", "p2", 0., 0.0, 1., 0.35 )
    
    p1.SetBottomMargin(1E-5)
    p2.SetTopMargin(1E-5)
    p1.Draw()
    p2.Draw()

    p1.cd()
    upper_1.SetLineColor(kBlue)
    upper_1.SetMarkerColor(kBlue)
    upper_1.SetMarkerStyle(kFullCircle)
    
    upper_1.Draw("E1")

    upper_1.GetYaxis().SetTitle("Backscattering Fraction")
    upper_1.GetXaxis().SetRangeUser(00.0, 5100)
    upper_1.GetYaxis().SetRangeUser(0.001, 0.089)
#    FormatAxis(upper_1, xtsize = 0., xlsize = 0)
    
    lower_1.SetMarkerStyle(kFullSquare)
    lower_1.SetMarkerColor(kBlack)
    lower_1.SetLineColor(kBlack)
    lower_1.Draw("E1 SAME")

    upper_2.SetMarkerStyle(kOpenCircle)
    upper_2.SetMarkerColor(kBlue)
    upper_2.SetLineColor(kBlue)
    upper_2.Draw("E1X0 SAME")
    
    lower_2.SetMarkerStyle(kOpenSquare)
    lower_2.SetMarkerColor(kBlack)
    lower_2.SetLineColor(kBlack)
    lower_2.Draw("E1X0 SAME")

    # upper_resid = TGraphError
    # lower_resid = lower_1.Clone()
    chi2_l = 0.
    chi2_u = 0.
    ndf_u = 0
    ndf_l = 0
    x_u = array('d', [])
    y_u = array('d', [])
    x_l = array('d', [])
    y_l = array('d', [])

    for i in range(1, upper_1.GetNbinsX()+1) :
        if (upper_1.GetBinError(i) > 0 and upper_2.GetBinError(i) > 0 ) :
#            upper_1.GetBinCenter(i) > 1000.) :
#            print upper_1.GetBinCenter(i)
            ty = (upper_1.GetBinContent(i) - upper_2.GetBinContent(i))
            er = sqrt((upper_1.GetBinError(i))**2 + (upper_2.GetBinError(i))**2)
            ty = ty / er

            x_u.append(upper_1.GetBinCenter(i))
            y_u.append(ty)
            
            if (upper_1.GetBinCenter(i) > 1000) :
                print ty
                chi2_u = chi2_u + ty**2
                ndf_u = ndf_u + 1

        if (lower_1.GetBinError(i) > 0 and lower_2.GetBinError(i) > 0) :
#            lower_1.GetBinCenter(i) > 1000.) :
            ty = (lower_1.GetBinContent(i) - lower_2.GetBinContent(i))
            er = sqrt((lower_1.GetBinError(i))**2 + (lower_2.GetBinError(i))**2)
            ty = ty / er

            x_l.append(lower_1.GetBinCenter(i))
            y_l.append(ty)

            if (lower_1.GetBinCenter(i) > 1000) :
                chi2_l = chi2_l + ty**2
                ndf_l = ndf_l + 1

    upper_resid = TGraph(len(x_u), x_u, y_u)
    lower_resid = TGraph(len(y_l), x_l, y_l)
    E = np.linspace(0, 5.1, 100)
    aE = array('d', [0])
    zeta = array('d', [0])
    zeta_up = array('d', [0])
    zeta_lo = array('d', [0])
    for i in range(len(E)) :
        if (E[i] > 0.05) :
            if (len(aE) == 1 and aE[0] == 0) :
                aE[0] = E[i]*1000
                z = Backscattering_coefficient(6., E[i])
                zeta[0] = z.nominal_value
                zeta_up[0] = z.nominal_value + z.std_dev
                zeta_lo[0] = z.nominal_value - z.std_dev
            else :
                aE.append(E[i]*1000)
                z = Backscattering_coefficient(6., E[i])
                zeta.append(z.nominal_value)
                zeta_up.append(z.nominal_value + z.std_dev)
                zeta_lo.append(z.nominal_value - z.std_dev)

    # g = TGraph(len(aE), aE, zeta)
    # g.SetTitle("Empirical (Tabata1971)")
    #g.Draw("LSAME")
    gmi = TGraph(len(aE), aE, zeta)
    gup = TGraph(len(aE), aE, zeta_up)
    glo = TGraph(len(aE), aE, zeta_lo)
    for g in [gup, glo] :
       g.SetLineColor(2)
       g.SetLineStyle(2)
       g.SetLineWidth(2)
       g.Draw("LSAME")
       g.SetTitle("Tabata1971")

    gPad.BuildLegend()
    gPad.Update()

    for h in [upper_resid, lower_resid] :
        h.GetXaxis().SetRangeUser(p1.GetUxmin(), p1.GetUxmax())
        h.GetYaxis().SetRangeUser(-4., 4.)
        h.GetYaxis().SetTitle("(Data - Sim) / #sigma")
        h.GetYaxis().SetTitleSize(0.11)
        h.GetYaxis().SetTitleOffset(0.4)
        h.GetYaxis().SetLabelSize(0.10)
        h.GetXaxis().SetTitleSize(0.08)
        h.GetXaxis().SetLabelSize(0.08)
        h.SetMarkerColor(mycol[0])
        h.SetMarkerStyle(kFullCircle)
        h.GetYaxis().SetLabelOffset(0.015)
        
    p2.cd()
#    upper_resid.Draw("AP")
    upper_resid.RemovePoint(0)
#    upper_resid.RemovePoint(0)
    lower_resid.RemovePoint(0)
    gPad.Update()
    l1 = TF1("l1", "1", p1.GetUxmin(), p1.GetUxmax())
    l2 = TF1("l2", "-1",p1.GetUxmin(), p1.GetUxmax())
    l3 = TF1("l3", "3", p1.GetUxmin(), p1.GetUxmax())
    l4 = TF1("l4", "-3",p1.GetUxmin(), p1.GetUxmax())

    l1.Draw("L")
    l2.Draw("L SAME")
    l3.Draw("L SAME")
    l4.Draw("L SAME")
    gPad.Update()
    upper_resid.Draw("P SAME")
    l1.SetLineStyle(2)
    l2.SetLineStyle(2)

    gPad.Update()

#    for l in [l1, l2, l3, l4] :
#        l.SetLineWidth(2)
#        l.Draw("L S")

    
    lower_resid.SetMarkerColor(kBlack)
    lower_resid.SetMarkerStyle(kFullSquare)
    lower_resid.Draw("P SAME")
    gPad.Update()
    l5 = TF1("l5", "1", gPad.GetUxmin(), gPad.GetUxmax())
    l6 = TF1("l6", "-1", gPad.GetUxmin(), gPad.GetUxmax())
    l5.SetLineStyle(2)
    l6.SetLineStyle(2)
    l7 = TF1("l7", "3", gPad.GetUxmin(), gPad.GetUxmax())
    l8 = TF1("l8", "-3",  gPad.GetUxmin(), gPad.GetUxmax())

    for l in [l5, l6, l7, l8] :
        l.SetLineWidth(2)
        #l.Draw("L SAME")
    print "upper chi2/n = ", chi2_u, " / ", ndf_u
    print "lower chi2/n = ", chi2_l, " / ", ndf_l

    raw_input("Enter to end")

    return [gup, glo]
