import numpy
from math import sqrt
from ROOT import *

#ROOT .Add method is buggy
def add_histos(h1,h2,a):
    hdiff = h1.Clone()
    hdiff.Reset()
    nx, ny = h1.GetNbinsX(), h1.GetNbinsY()
    for ibx in range(nx):
        for iby in range(ny):
            c = (h1.GetBinContent(ibx,iby)
                 +float(a)*h2.GetBinContent(ibx,iby))
            hdiff.SetBinContent(ibx,iby,c)
    return hdiff.Clone()

#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

#returns a histogram for only the fiducial
def fid_only(h1,fid,pcut_min=22.999,pcut_max=27.001):
    h2 = h1.Clone()
    h2.Reset()
    nx, ny = h1.GetNbinsX(), h1.GetNbinsY()
    for ibx in range(nx):
        p = h1.GetXaxis().GetBinCenter(ibx)
#        p = h1.GetXaxis().GetBinCenter(ibx+1)
#        p = h1.GetXaxis().GetBinCenter(ibx-1)
        if pcut_min < p < pcut_max: continue
        for iby in range(ny):
            ct = h1.GetYaxis().GetBinCenter(iby)
#            ct = h1.GetYaxis().GetBinCenter(iby+1)
#            ct = h1.GetYaxis().GetBinCenter(iby-1)
            if in_fid(p,ct,fid):
                h2.SetBinContent(ibx,iby,h1.GetBinContent(ibx,iby))
    return h2.Clone()

#make simple histogram of inefficiencies
def simple_histo(h1):
    h = TH1D('simple','simple',100,0,0.01)
    h.GetXaxis().SetNdivisions(505)
    for ibx in range(h1.GetNbinsX()):
        for iby in range(h1.GetNbinsY()):
            c = h1.GetBinContent(ibx,iby)
            if 0.<c<1.: h.Fill(c)
    print 'mean=',h.GetMean(),'; RMS=',h.GetRMS()
    return h.Clone()

#take x projection
def project_vs_p(h1):
    h2 = h1.ProjectionX().Clone()
    h2.Reset()
    for ibx in range(h1.GetNbinsX()):
        x = []
        for iby in range(h1.GetNbinsY()):
            c = h1.GetBinContent(ibx,iby)
            if 0.<c<1.: x.append(c)
        if len(x)>0:
            h2.SetBinContent(ibx,sum(x)/len(x))
#            h2.SetBinError(ibx,numpy.std(x)/sqrt(len(x)-1))
            h2.SetBinError(ibx,numpy.std(x)/sqrt(len(x)))
    h2.SetMarkerStyle(20)
    h2.SetMaximum(0.005)
    return h2.Clone()

#take ct projection
def project_vs_ct(h1):
    h2 = h1.ProjectionY().Clone()
    h2.Reset()
    for iby in range(h1.GetNbinsY()):
        x = []
        for ibx in range(h1.GetNbinsX()):
            c = h1.GetBinContent(ibx,iby)
            if 0.<c<1.: x.append(c)
        if len(x)>0:
            h2.SetBinContent(iby,sum(x)/len(x))
#            h2.SetBinError(iby,numpy.std(x)/sqrt(len(x)-1))
            h2.SetBinError(iby,numpy.std(x)/sqrt(len(x)))
    h2.SetMarkerStyle(20)
    h2.SetMaximum(0.005)
    return h2.Clone()