import ROOT,math
from array import array
#from ROOT import std,TMath

colors = [ROOT.kBlack,ROOT.kRed+1,ROOT.kGreen+3,ROOT.kBlue+1,ROOT.kMagenta+1,ROOT.kOrange+1,ROOT.kViolet-7]

def calcMeanErr2(s1,ds1,h1,dh1,s2,ds2,h2,dh2) :

    divider=(h1*s1+h2*s2)
    divider *= divider

    ddh1  = math.fabs(h2*s1*s1*s2-h2*s2*s2*s1)
    ddh1 /= divider
    ddh1 *= dh1

    dds1  = math.fabs(h1*h1*s1*s1+2*h1*h2*s1*s2-h2*h1*s2*s2)
    dds1 /=divider
    dds1 *= ds1

    ddh2  = math.fabs(h1*s1*s2*s2-h1*s1*s1*s2*s2)
    ddh2 /= divider
    ddh2 *= dh2

    dds2  = math.fabs(h2*h2*s2*s2+2*h2*h1*s1*s2 - h1*h2*s1*s1)
    dds2 /= divider
    dds2 *= ds2

    meanerr = math.sqrt(ddh1+ddh2+dds1+dds2)

    return meanerr


def fitres(h,df,factors=[1.,1.,2.,12.]) :
#    histo.SetStats(ROOT.kFALSE) 
    data={}
    #    print '***************i=',i,colors[i],'*********************'
    #histo.SetLineColor(colors[i])
    histo=h.Clone()
    histo.SetFillColor(ROOT.kAzure-8)
    testfit = ROOT.TF1("testfitY"+str(1),"gaus",-10000.,10000.)
    testfit.SetLineColor(2)
    histo.Fit(testfit, "QN+", "", -10000,10000)

    rfit = ROOT.TF1("fitfuncY"+str(1),"gaus + gaus(3)",-10000.,10000.)
    rfit1 = ROOT.TF1("gaus1","gaus",-10000.,10000.)
    rfit2 = ROOT.TF1("gaus2","gaus",-10000.,10000.)

    rfit.SetNpx(1000)
    rfit.SetParName(0,"const1")
    rfit.SetParName(1,"mean1")
    rfit.SetParName(2,"sigma1")
    rfit.SetParName(3,"const2")
    rfit.SetParName(4,"mean2")
    rfit.SetParName(5,"sigma2")
     
    tconst = [None]*2
    tmean = [None]*2
    tsigma = [None]*2

    tconst[0] = tconst[1] = testfit.GetParameter(0)
    tmean[0]  = tmean[1]  = 0
    tsigma[0] = tsigma[1] = testfit.GetParameter(2)

    #medium gaus
    rfit.SetParameter(0,tconst[0]/1.5)
    rfit.SetParameter(1,0.)
    #    rfit.SetParLimits(1,0.,0.)
    #    rfit.FixParameter(1,0.)
    rfit.SetParameter(2,tsigma[0]*2.5)
    rfit.SetParLimits(2,tsigma[0]*factors[2], tsigma[0]*factors[3])
    
    #narrow gaus
    rfit.SetParameter(3, tconst[1])
    rfit.SetParameter(4, 0.)
    #    rfit.FixParameter(4,0.)
    rfit.SetParameter(5, tsigma[1]*0.5)
    rfit.SetParLimits(5, 0.0*factors[0], tsigma[1]*factors[1])
  
    rfit.SetLineWidth(1)
    print "now fitting in ranges:",-3*histo.GetRMS(),3*histo.GetRMS()
    histo.Fit(rfit, "MIQ+","", -3*histo.GetRMS(),3*histo.GetRMS())

    if rfit.GetNDF()>0:
        scaleu=rfit.GetChisquare()/rfit.GetNDF()
        scaled=1/scaleu
    else:
        scaleu=2.
        scaled=0.5

    if df==1:
        for j in range(2) :
            tconst[j]=rfit.GetParameter(j*2)
            tmean[j]=rfit.GetParameter(j*2+1)
            tsigma[j]=rfit.GetParameter(j*2+2)

        #        narrow gaus
        rfit.SetParameter(0,tconst[0])
        #    rfit.SetParLimits(0,tconst[0]*0.5, tconst[0]*2)
        rfit.SetParameter(2,tsigma[0])
        rfit.SetParLimits(2,tsigma[0]*scaled, tsigma[0]*scaleu)
    
        #        broadest gaus
        rfit.SetParameter(3, tconst[1])
        #    rfit.SetParLimits(3, 0, tconst[1]/1.5)
        rfit.SetParameter(5, tsigma[1])
        rfit.SetParLimits(5, tsigma[1]*scaled, tsigma[1]*scaleu)
    
        rfit.SetLineWidth(2)
        histo.Fit(rfit, "+", "", -10000,10000)


    if 0:
#        canv=ROOT.TCanvas("bla","bla")
#        leg=ROOT.TLegend(0.7,.7,0.9,.9,"","NDC")
#        canv.cd()
#        histo.Draw()
        rfit1.SetParameter(0,rfit.GetParameter(0))
        rfit1.SetParameter(1,rfit.GetParameter(1))
        rfit1.SetParameter(2,rfit.GetParameter(2))
        rfit1.SetLineColor(ROOT.kRed)
#        leg.AddEntry(rfit1,"Broad")
        rfit1.DrawCopy("same")
        rfit2.SetParameter(0,rfit.GetParameter(3))
        rfit2.SetParameter(1,rfit.GetParameter(4))
        rfit2.SetParameter(2,rfit.GetParameter(5))
        rfit2.SetLineColor(ROOT.kBlue)
#        leg.AddEntry(rfit2,"Narrow")
        rfit2.DrawCopy("same")
#        testfit.SetLineColor(ROOT.kGreen)
#        testfit.Draw("same")
#        leg.AddEntry(testfit,"testfit")
#        leg.Draw()
#        print "*******************************************"
#        print tsigma[0]*10.2, tsigma[1]*0.05, tsigma[1]*1.5
#        print "*******************************************"
#        drawtext(["broad"+str(rfit.GetParameter(2)),"small"+str(rfit.GetParameter(5))])
#        nix=raw_input("blahorst")


        
    int1 = rfit.GetParameter(0)*rfit.GetParameter(2) #narrow gaus
    int2 = rfit.GetParameter(3)*rfit.GetParameter(5) #broad gaus


    meanerr = 0
    if int1 > 0 :
        ratio = int2/int1

    #mean resolution:
 
        data['csig'] = (int1*rfit.GetParameter(2)+int2*rfit.GetParameter(5))/(int1+int2) 
        data['csigerr'] = calcMeanErr2(rfit.GetParameter(2), rfit.GetParError(2),
                                       rfit.GetParameter(0), rfit.GetParError(0),
                                       rfit.GetParameter(5), rfit.GetParError(5),
                                       rfit.GetParameter(3), rfit.GetParError(3))
            
#            print "***********************************"
#            print "small=",rfit.GetParameter(5)," medium=",rfit.GetParameter(2)," mean=",meanRes
#            print "tsigma",tsigma," tmean",tmean," tconst",tconst
#            print "***********************************"

            
    else:
        data['csig'] = -9999
        data['csigerr']=0.


    data['nsig']    = rfit.GetParameter(5)
    data['nsigerr'] = rfit.GetParError(5)
    data['bconst']  = rfit.GetParameter(0)
    data['nconst']  = rfit.GetParameter(3)
    data['bmean']   = rfit.GetParameter(1)
    data['nmean']   = rfit.GetParameter(4)
        
    if rfit.GetNDF()>0:
        data['chindf']=rfit.GetChisquare()/rfit.GetNDF()
    else:
        data['chindf']=0

    data['bsig']=rfit.GetParameter(2)
    data['bsigerr']=rfit.GetParError(2)
    return data

def drawresdata(data,opt=""):
    text=ROOT.TLatex()
    text.SetTextColor(ROOT.kBlack)
    text.SetTextSize(0.05)
    text.SetNDC()
    y=0.8
    ystep=0.05
    text.DrawLatex(0.15,y,"#sigma_{narrow}= "+str(round(data['nsig'],1))+"#pm"+str(round(data['nsigerr'],1))+" #mum")
    y-=ystep
    text.DrawLatex(0.15,y,"#sigma_{broad}="+str(round(data['bsig'],1))+"#pm"+str(round(data['bsigerr'],1))+" #mum")
    y-=ystep
    text.DrawLatex(0.15,y,"#sigma_{mean}= "+str(round(data['csig'],1))+"#pm"+str(round(data['csigerr'],1))+" #mum")
    if opt=="mean":
        y-=ystep
        text.DrawLatex(0.15,y,"Mean= "+str(round(data['nmean'],3)))
    y-=ystep
    text.DrawLatex(0.15,y,"#chi^{2}/ndf="+str(round(data['chindf'],1)))
    y-=ystep
    if (data['bsig']*data['bconst'])>0:
        text.DrawLatex(0.15,y,"S/B="+str(round((data['nsig']*data['nconst'])/(data['bsig']*data['bconst']),2)))
    
    return

def drawGaussians(data):
    rfit1 = ROOT.TF1("gaus1","gaus",-10000.,10000.)
    rfit1.SetParameter(0,data['bconst'])
    rfit1.SetParameter(1,data['bmean'])
    rfit1.SetParameter(2,data['bsig'])
    rfit1.SetLineColor(ROOT.kRed)
    rfit1.DrawCopy("same")
   
    rfit2 = ROOT.TF1("gaus2","gaus",-10000.,10000.) 
    rfit2.SetParameter(0,data['nconst'])
    rfit2.SetParameter(1,data['nmean'])
    rfit2.SetParameter(2,data['nsig'])
    rfit2.SetLineColor(ROOT.kBlue)
    rfit2.DrawCopy("same")
    
    return

def drawPrelim() :
    prelim=ROOT.TLatex()
    prelim.SetTextColor(ROOT.kGray)
    prelim.SetTextSize(0.07)
    prelim.SetTextAngle(20)
    prelim.SetNDC()
    prelim.DrawLatex(0.38,0.45,"preliminary")
    return

def drawDiff(noB,ar,ne):
    if noB==1:
        if ar==1:
            DIFFX = 0.0227*10000   #(mu m/ mu s) without B-Field
    else:
        if ar==1:
            DIFFX = 0.020534*10000 #(mu m/ mu s)Argon with B-Field
        if ne==1:
            DIFFX = 0.019833*10000 #(mu m/ mu s)Neon with B-Field

    diffT = ROOT.TF1("diffT","[0]*TMath::Sqrt(x)",0,75)
    diffT.SetNpx(1000)
    diffT.SetLineWidth(2)
    diffT.SetLineStyle(4)
    diffT.SetParameter(0,DIFFX)
    diffT.SetTitle("")

    return diffT

def sort_inner(inner):
#inner is each inner list in the list of lists to be sorted
#(here item at index 1 of each inner list is to be sorted)
    return inner[0]

def sort_inner2(inner):
    return inner[2]

def Drawh(hist,hists,legar,canv):
    
    col=0
    colors = [ROOT.kBlack,ROOT.kRed+1,ROOT.kRed-7,ROOT.kGreen+3,ROOT.kBlue+1,ROOT.kMagenta+1,ROOT.kOrange+1,ROOT.kViolet-7]
    leg=ROOT.TLegend(0.7,.7,0.9,.9,"","NDC")
    canv.cd()
    maxbin=0
    for h in hist:
        col+=1
        if col < len(colors):
            color=colors[col]
        else:
            print "no color available"
            color=1
#        hist[h].SetTitle(hists)
        if hist[h].GetMaximum() > maxbin:
            maxbin=hist[h].GetMaximum()
        hist[h].SetLineColor(color)
        if col == 1:
            hist[h].GetXaxis().SetTitle(axlabel[0].get(hists,""))
            hist[h].GetYaxis().SetTitle(axlabel[1].get(hists,""))
            hist[h].Draw()
            first=hist[h]
        else:
            hist[h].Draw("same")
        leg.AddEntry(hist[h],h)
    first.SetMaximum(maxbin*1.02)
    leg.SetFillColor(0)
    leg.Draw()
    legar.append(leg)
#    else:
#        hist.Draw()

    return 

def Divideh(h1,h2,name,title):
    temphist=h1.Clone(name)
    temphist.SetTitle(title)
    temphist.Divide(h2)
    return temphist

def Addh(h1,h2,factor,name,title=""):
    temphist=h1.Clone(name)
    temphist.SetTitle(title)
    temphist.Add(h2,factor)
    return temphist

def fillcutred(cut,event,hist):
#    hist=histsdict['cut reduction']
    hist.Fill(1,cut)

    if event['clmeanamp']<400:
        hist.Fill(2,cut)
    if event['theta']<10 or event['theta']>170:
        hist.Fill(3,cut)
    if abs(event['phi'])<2:
        hist.Fill(4,cut)
    if event['clmeanamp']>400 and (event['theta']<10 or event['theta']>170):
        hist.Fill(5,cut)
    return

def plotfullevent(thevent):
    trackcounter=0
    digamp=[]
    digpos=[]
    canvamp=[]
    canvpos=[]
    lines=[]
    lines2d=[]
    lines2d2=[]
    for tr in range(len(thevent)):
        trackcounter+=1
        vec=ROOT.TVector3(1.,1.,1.)
        vec.SetPhi(thevent[tr]['phi'])
        vec.SetTheta(thevent[tr]['theta'])
        vec.SetMag(20.)
        line=ROOT.TPolyLine3D(2,"")
        x2=thevent[tr]['hits'][len(thevent[tr]['hits'])-1]['pos'][0]
        y2=thevent[tr]['hits'][len(thevent[tr]['hits'])-1]['pos'][1]
        z2=thevent[tr]['hits'][len(thevent[tr]['hits'])-1]['pos'][2]
        x1=thevent[tr]['hits'][0]['pos'][0]
        y1=thevent[tr]['hits'][0]['pos'][1]
        z1=thevent[tr]['hits'][0]['pos'][2]
        line.SetPoint(0,x1,y1,z1)
        line.SetPoint(1,x2,y2,z2)
        lines2d.append(ROOT.TLine(x1,y1,x2,y2))
        lines2d2.append(ROOT.TLine(z1,y1,z2,y2))
        lines2d[trackcounter-1].SetLineColor(trackcounter)
        lines2d2[trackcounter-1].SetLineColor(trackcounter)

        print "appending line, size:", len(lines2d), len(lines2d2)                    
                    #tcn2.cd()
                    #hits3d.Draw()
                    #vec.Draw("same")
                    #line.Draw("same");
                    #tcn2.Update()

        canvamp.append(ROOT.TCanvas("samples"+str(trackcounter),"samples"+str(trackcounter),720,10,700,500))
        canvpos.append(ROOT.TCanvas("samples2D"+str(trackcounter),"samples2D"+str(trackcounter),720,520,700,500))
        print int(math.ceil(math.sqrt(len(thevent[tr]['hits']))))
        canvamp[trackcounter-1].Divide(int(math.ceil(math.sqrt(len(thevent[tr]['hits'])))),int(math.ceil(math.sqrt(len(thevent[tr]['hits'])))))
        canvpos[trackcounter-1].Divide(int(math.ceil(math.sqrt(len(thevent[tr]['hits'])))),int(math.ceil(math.sqrt(len(thevent[tr]['hits'])))))
        clcount=1
        hcount=0
#        his=[]
#        his2d=[]
        digpos.append([])
        digamp.append([])
        for clust in thevent[tr]['hits']:
            histsdict['hxy'].Fill(clust['pos'][0],clust['pos'][1])#,clust['res'][0])
            histsdict['hyz'].Fill(clust['pos'][2],clust['pos'][1])#,clust['res'][0])
            canvamp[trackcounter-1].cd(clcount)
            clcount+=1
            same=""
            i=1
            digpos[trackcounter-1].append(ROOT.TH2D("cl2d"+str(clcount-1)+"_"+str(hcount),"cluster "+str(clcount-1),200,-16.,16.,200,-16.,16.))

            for dig in clust['digis']:
                digamp[trackcounter-1].append(ROOT.TH1D("cl"+str(clcount-1)+"_"+str(hcount),"cluster "+str(clcount-1),511,0.,511.))

                digamp[trackcounter-1][hcount].Fill(dig['time'],dig['amp'])
                digpos[trackcounter-1][clcount-2].Fill(mapping[int(dig['pad'])][0],mapping[int(dig['pad'])][1])

                digamp[trackcounter-1][hcount].SetLineColor(i)
                digamp[trackcounter-1][hcount].Draw(same)
                same="same"
                hcount+=1
                i+=1
                print "pads:",int(dig['pad']), mapping[int(dig['pad'])][0],mapping[int(dig['pad'])][1]
            canvpos[trackcounter-1].cd(clcount-1)
            digpos[trackcounter-1][clcount-2].Draw("colz")
            print ""

        print "phi angle:",thevent[tr]['phi']
        print "theta angle:",thevent[tr]['theta']
        print "momentum:",thevent[tr]['mom']
        print "chi2",thevent[tr]['chi2']," ndf",thevent[tr]['ndf']
        print "clmeanamp:",thevent[tr]['clmeanamp']
        print "event number:", evt
        print "vector",vec.X(),vec.Y(),vec.Z()
        print "cluster1",thevent[tr]['hits'][0]['pos'][0],thevent[tr]['hits'][0]['pos'][1],thevent[tr]['hits'][0]['pos'][2]
        print "cluster2",clust['pos'][0],clust['pos'][1],clust['pos'][2]
        print "trackcounter: ",trackcounter

        tcn.cd(1)
        histsdict['hxy'].Draw("colz")
        tcn.cd(2)
        histsdict['hyz'].Draw("colz")
        tcn.Update()

        canvamp[trackcounter-1].Update()
        canvpos[trackcounter-1].Update()
#        if trackcounter==len(event):
    for l in range(len(lines2d)):
        tcn.cd(1)
        lines2d[l].Draw("same")
        tcn.cd(2)
        lines2d2[l].Draw("same")
        tcn.Update()                            
    raw_input("next event")

#    tcn3.Delete()
#    tcn4.Delete()
#    for i in range(len(his)):
#        his[i].Delete()
#    del his
    histsdict['hxy'].Reset()
    histsdict['hyz'].Reset()

    return

def parseRuns(theruns):
    runList=[]
    colindex = theruns.find(",")
    if colindex > 0:
        runs = theruns.split(",")
        for i in range(len(runs)) :
            dashIndex = runs[i].find("-")
            if dashIndex > 0 :
                runs[i] = runs[i].split("-")
                for j in range(int(runs[i][0]), int(runs[i][1])+1) :
                    runList.append(j)
            if dashIndex < 0 and len(runs) > 1 :
                runList.append(int(runs[i]))
    else :
        dashIndex = theruns.find("-")
        if dashIndex > 0 :
            runs = theruns.split("-")
            runList = range(int(runs[0]), int(runs[1])+1)
        else:
            runList = [theruns]
    return runList 

def set_palette(name="palette", ncontours=999):
    """Set a color palette from a given RGB list
    stops, red, green and blue should all be lists of the same length
    see set_decent_colors for an example"""

    if name == "gray" or name == "grayscale":
        stops = [0.00, 0.34, 0.61, 0.84, 1.00]
        red   = [1.00, 0.84, 0.61, 0.34, 0.00]
        green = [1.00, 0.84, 0.61, 0.34, 0.00]
        blue  = [1.00, 0.84, 0.61, 0.34, 0.00]
    # elif name == "whatever":
        # (define more palettes)
    else:
        # default palette, looks cool
        stops = [0.00, 0.34, 0.61, 0.84, 1.00]
        red   = [0.00, 0.00, 0.87, 1.00, 0.51]
        green = [0.00, 0.81, 1.00, 0.20, 0.00]
        blue  = [0.51, 1.00, 0.12, 0.00, 0.00]

    s = array('d', stops)
    r = array('d', red)
    g = array('d', green)
    b = array('d', blue)

    npoints = len(s)
    ROOT.TColor.CreateGradientColorTable(npoints, s, r, g, b, ncontours)
    ROOT.gStyle.SetNumberContours(ncontours)

def setRange(h,mini=0,maxi=0):
    for i in range(h.GetNbinsX()+1):
        for j in range(h.GetNbinsY()+1):
            if (maxi!=mini):
                cont=h.GetBinContent(i,j)
                if (cont>maxi):
                    h.SetBinContent(i,j,maxi)
                if (cont<mini):
                    h.SetBinContent(i,j,mini)
    if(mini!=maxi):
        h.GetZaxis().SetRangeUser(mini,maxi)
    return

def createGraph(hists,points,toplot="mean",color=1,style=20):
    graph=ROOT.TGraph()
    counter=-1
    mini=100000
    maxi=0
    for h in hists:
        counter+=1
        val=5
        if toplot=="mean":
            val=h.GetMean()
        if toplot=="int2":
            nb=h.GetNbinsX()
            val=0
            for i in range(1,nb):
                val+=h.GetBinContent(i)*h.GetXaxis().GetBinCenter(i)
        maxi=max(val,maxi)
        mini=min(val,mini)
        graph.SetPoint(counter,points[counter],val)

    graph.SetMarkerStyle(style)
    graph.SetMarkerColor(color)
    graph.SetMaximum(maxi)
    graph.SetMinimum(mini)
    return graph

def setGraphLayout(g,xtit,ytit,marker=20,size=1,col=1):

    return g