import yoda
import ROOT
from array import array


def _TH1toE1D(th1):

    edges = [ ]
    for ix in range(th1.GetNbinsX()):
        axx = th1.GetXaxis()
        if not ix:  edges.append(axx.GetBinLowEdge(ix+1))
        edges.append(axx.GetBinUpEdge(ix+1))

    rtn = yoda.Estimate1D(edges)
    for ix in range(1,th1.GetNbinsX()+1):
        rtn.bin(ix).setVal(th1.GetBinContent(ix))
        rtn.bin(ix).setErr(th1.GetBinError(ix))

    rtn.setPath(th1.GetName())
    rtn.setTitle(th1.GetTitle())
    rtn.setAnnotation("XLabel", th1.GetXaxis().GetTitle())
    rtn.setAnnotation("YLabel", th1.GetYaxis().GetTitle())

    return rtn;


def _TH2toE2D(th2):

    xedges = [ ]
    yedges = [ ]
    for ix in range(th2.GetNbinsX()):
        axx = th2.GetXaxis()
        if not ix:  xedges.append(axx.GetBinLowEdge(ix+1))
        xedges.append(axx.GetBinUpEdge(ix+1))

    for iy in range(th2.GetNbinsY()):
        axy = th2.GetYaxis()
        if not iy:  yedges.append(axy.GetBinLowEdge(iy+1))
        yedges.append(axy.GetBinUpEdge(iy+1))

    rtn = yoda.Estimate2D(xedges,yedges)
    for iy in range(1,th2.GetNbinsY()+1):
        for ix in range(1,th2.GetNbinsX()+1):
            rtn.bin(ix,iy).setVal(th2.GetBinContent(ix,iy))
            rtn.bin(ix,iy).setErr(th2.GetBinError(ix,iy))

    rtn.setPath(th2.GetName())
    rtn.setTitle(th2.GetTitle())
    rtn.setAnnotation("XLabel", th2.GetXaxis().GetTitle())
    rtn.setAnnotation("YLabel", th2.GetYaxis().GetTitle())
    rtn.setAnnotation("ZLabel", th2.GetZaxis().GetTitle())

    return rtn;


def _TH1toS2D(th1, widthscale):
    rtn = yoda.Scatter2D()
    for i in range(1, th1.GetNbinsX()+1):
        axx = th1.GetXaxis()
        x = axx.GetBinCenter(ix)
        exminus = x - axx.GetBinLowEdge(ix)
        explus = axx.GetBinUpEdge(ix) - x
        xwidth = axx.GetBinWidth(ix) if widthscale else 1.0

        val = th1.GetBinContent(ix) / xwidth
        evalminus = th1.GetBinErrorLow(ix) / xwidth
        evalplus = th1.GetBinErrorUp(ix) / xwidth
        rtn.addPoint((x, val), (exminus, evalminus), (explus, evalplus))

    rtn.setPath(th1.GetName())
    rtn.setTitle(th1.GetTitle())
    rtn.setAnnotation("XLabel", th1.GetXaxis().GetTitle())
    rtn.setAnnotation("YLabel", th1.GetYaxis().GetTitle())

    return rtn


def _TH2toS3D(th2, widthscale):

    rtn = yoda.Scatter3D()
    for ix in range(1, th2.GetNbinsX()+1):
        for iy in range(1, th2.GetNbinsY()+1):
            axx = th2.GetXaxis()
            x = axx.GetBinCenter(ix)
            exminus = x - axx.GetBinLowEdge(ix)
            explus = axx.GetBinUpEdge(ix) - x
            xwidth = axx.GetBinWidth(ix)

            axy = th2.GetYaxis()
            y = axy.GetBinCenter(iy)
            eyminus = y - axy.GetBinLowEdge(iy)
            eyplus = axy.GetBinUpEdge(iy) - y
            ywidth = axy.GetBinWidth(iy)

            ixy = th2.GetBin(ix, iy)
            area = xwidth*ywidth if widthscale else 1.0
            val = th2.GetBinContent(ixy) / area
            evalminus = th2.GetBinErrorLow(ixy) / area
            evalplus = th2.GetBinErrorUp(ixy) / area
            rtn.addPoint((x, y), (exminus, evalminus), (explus, evalplus))

    rtn.setPath(th2.GetName())
    rtn.setTitle(th2.GetTitle())
    rtn.setAnnotation("XLabel", th2.GetXaxis().GetTitle())
    rtn.setAnnotation("YLabel", th2.GetYaxis().GetTitle())
    rtn.setAnnotation("ZLabel", th2.GetZaxis().GetTitle())

    return rtn;



def to_yoda(root_obj, asScatter=False, widthscale=False):
    if isinstance(root_obj, ROOT.TH2):
        if asScatter:
            return _TH2toS3D(root_obj, widthscale)
        else:
            return _TH2toE2D(root_obj)
    elif isinstance(root_obj, ROOT.TProfile):
        if asScatter:
            return _TH1toS2D(root_obj, widthscale)
        else:
            return _TH1toE1D(root_obj)
    elif isinstance(root_obj, ROOT.TH1):
        if asScatter:
            return _TH1toS2D(root_obj, widthscale)
        else:
            return _TH1toE1D(root_obj)




def _H1toTH1D(h1d, widthscale):
    # Book ROOT histogram
    edges = h1d.xEdges()
    try:
        edgesarr = array('d', edges)
    except TypeError as te:
        raise Exception(f"Edge error when converting Histo1D: {te}")
    rtn = ROOT.TH1D(h1d.path(), h1d.title(), len(edges)-1, edgesarr);
    rtn.Sumw2()
    sumw2s = rtn.GetSumw2()
    for i in range(rtn.GetNbinsX()+2):
        try:
            b = h1d.bin(i)
            rtn.SetBinContent(i, b.sumW() / (b.dVol() if widthscale else 1.0))
            sumw2s[i] = b.sumW2()
        except:
            pass
    # Labels
    if h1d.hasAnnotation("XLabel"): rtn.SetXTitle(h.annotation("XLabel"))
    if h1d.hasAnnotation("YLabel"): rtn.SetYTitle(h.annotation("YLabel"))
    return rtn


def _P1toTProfile(p1d):
    # Book ROOT histogram
    edges = p1d.xEdges()
    try:
        edgesarr = array('d', edges)
    except TypeError as te:
        raise Exception(f"Edge error when converting Profile1D: {te}")
    rtn = ROOT.TProfile(p1d.path(), p1d.title(), len(edges)-1, edgesarr);
    rtn.Sumw2()
    sumwys = rtn.GetArray()
    sumwy2s = rtn.GetSumw2()
    for i in range(rtn.GetNbinsX()+2):
        try:
            b = p1d.binAt(rtn.GetBinCenter(i))
            rtn.SetBinEntries(i, b.sumW())
            sumwys[i] = b.sumWY();
            sumwy2s[i] = b.sumWY2();
        except:
            pass
    # Labels
    if p1d.hasAnnotation("XLabel"):  rtn.SetXTitle(p1d.annotation("XLabel"))
    if p1d.hasAnnotation("YLabel"):  rtn.SetYTitle(p1d.annotation("YLabel"))
    return rtn


def _E1toTH1D(e1d):
    # Book ROOT histogram
    edges = e1d.xEdges()
    try:
        edgesarr = array('d', edges)
    except TypeError as te:
        raise Exception(f"Edge error when converting Estimate1D: {te}")
    rtn = ROOT.TH1D(e1d.path(), e1d.title(), len(edges)-1, edgesarr)
    rtn.Sumw2()
    sumw2s = rtn.GetSumw2()
    for i in range(rtn.GetNbinsX()+2):
      try :
        b = e1d.bin(i)
        rtn.SetBinContent(i, b.val())
        sumw2s[i] = b.totalErrAvg()*b.totalErrAvg();
      except:
        pass
    # Labels
    if e1d.hasAnnotation("XLabel"): rtn.SetXTitle(e1d.annotation("XLabel"))
    if e1d.hasAnnotation("YLabel"): rtn.SetYTitle(e1d.annotation("YLabel"))
    return rtn


def _H2toTH2D(h2d, widthscale):
    # Book ROOT histogram
    xedges = h2d.xEdges()
    yedges = h2d.yEdges()
    try:
        xedgesarr = array('d', xedges)
        yedgesarr = array('d', yedges)
    except TypeError as te:
        raise Exception(f"Edge error when converting Histo1D: {te}")
    rtn = ROOT.TH2D(h2d.path(), h2d.title(), len(xedges)-1, xedgesarr, len(yedges)-1, yedgesarr)
    rtn.Sumw2()
    sumw2s = rtn.GetSumw2()
    for ix in range(rtn.GetNbinsX()+2):
        for iy in range(rtn.GetNbinsY()+2):
            i = rtn.GetBin(ix, iy)
            try:
                b = h2d.binAt(rtn.GetXaxis().GetBinCenter(ix), rtn.GetYaxis().GetBinCenter(iy))
                rtn.SetBinContent(i, b.sumW() / (b.dVol() if widthscale else 1.0))
                sumw2s[i] = b.sumW2()
            except:
                pass
    # Labels
    if h2d.hasAnnotation("XLabel"): rtn.SetXTitle(h2d.annotation("XLabel"))
    if h2d.hasAnnotation("YLabel"): rtn.SetYTitle(h2d.annotation("YLabel"))
    if h2d.hasAnnotation("ZLabel"): rtn.SetZTitle(h2d.annotation("ZLabel"))
    return rtn


def _E2toTH2D(e2d):
    # Book ROOT histogram
    xedges = e2d.xEdges()
    yedges = e2d.yEdges()
    try:
        xedgesarr = array('d', xedges)
        yedgesarr = array('d', yedges)
    except TypeError as te:
        raise Exception(f"Edge error when converting Estimate2D: {te}")
    rtn = ROOT.TH2D(e2d.path(), e2d.title(), len(xedges)-1, xedgesarr, len(yedges)-1, yedges)
    rtn.Sumw2()
    sumw2s = rtn.GetSumw2()
    for ix in range(rtn.GetNbinsX()+2):
        for iy in range(rtn.GetNbinsY()+2):
            i = rtn.GetBin(ix, iy)
            try :
              #b = h.binAt(rtn.GetBinCenter(ix), rtn.GetBinCenter(iy)) // throws if in a gap
              rtn.SetBinContent(i, b.val())
              sumw2s[i] = b.totalErrAvg()*b.totalErrAvg()
            except:
                pass
    # Labels
    if e2d.hasAnnotation("XLabel"): rtn.SetXTitle(e2d.annotation("XLabel"))
    if e2d.hasAnnotation("YLabel"): rtn.SetYTitle(e2d.annotation("YLabel"))
    if e2d.hasAnnotation("ZLabel"): rtn.SetZTitle(e2d.annotation("ZLabel"))
    return rtn


def _S1toTGraph(s1d):
    xs, ys = [ ], [ ]
    exls, exhs = [ ], [ ]
    eyls, eyhs = [ ], [ ]
    for i in range(s1d.numPoints()):
        p = s1d.point(i)
        xs.append(i)
        ys.append(p.x())
        exls.append(0.0)
        exhs.append(0.0)
        eyls.append(p.xErrMinus())
        eyhs.append(p.xErrPlus())
    # Make the ROOT object... mm, the constructors don't take name+title, unlike all the histos!
    rtn = ROOT.TGraphAsymmErrors(len(xs), array('d', xs),   array('d', ys),
                                 array('d', exls), array('d', exhs),
                                 array('d', eyls), array('d', eyhs))
    rtn.SetName(s1d.path())
    rtn.SetTitle(s1d.title())
    # Labels
    if s1d.hasAnnotation("XLabel"): rtn.GetXaxis().SetTitle(s1d.annotation("XLabel"))
    if s1d.hasAnnotation("YLabel"): rtn.GetYaxis().SetTitle(s1d.annotation("YLabel"))
    return rtn


def _S2toTGraph(s2d):
    xs, ys = [ ], [ ]
    exls, exhs = [ ], [ ]
    eyls, eyhs = [ ], [ ]
    for i in range(s2d.numPoints()):
        p = s2d.point(i)
        xs.append(p.x())
        ys.append(p.y())
        exls.append(p.xErrMinus())
        exhs.append(p.xErrPlus())
        eyls.append(p.yErrMinus())
        eyhs.append(p.yErrPlus())
    # Make the ROOT object... mm, the constructors don't take name+title, unlike all the histos!
    rtn = ROOT.TGraphAsymmErrors(len(xs), array('d', xs),   array('d', ys),
                                 array('d', exls), array('d', exhs),
                                 array('d', eyls), array('d', eyhs))
    rtn.SetName(s2d.path())
    rtn.SetTitle(s2d.title())
    # Labels
    if s2d.hasAnnotation("XLabel"): rtn.GetXaxis().SetTitle(s2d.annotation("XLabel"))
    if s2d.hasAnnotation("YLabel"): rtn.GetYaxis().SetTitle(s2d.annotation("YLabel"))
    return rtn


def _H1toTGraph(h1d, binwidthdiv, usefocus, includeOverflows, includeHiddenBins):
    return _S2toTGraph(h1d.mkScatter(h1d.path(), binwidthdiv, usefocus, includeOverflows, includeMaskedBins))


def _P1toTGraph(p1d, binwidthdiv, usefocus, includeOverflows, includeHiddenBins):
    return _S2toTGraph(p1d.mkScatter(p1d.path(), binwidthdiv, usefocus, includeOverflows, includeMaskedBins))


def _E1toTGraph(e1d, includeOverflows, includeHiddenBins):
    return _S2toTGraph(e1d.mkScatter(e1d.path(), includeOverflows, includeMaskedBins))


def to_root(yoda_obj, asgraph=False, usefocus=False, widthscale=False,
                      includeOverflows=False, includeHiddenBins=False):
    if isinstance(yoda_obj, yoda.BinnedHisto1D):
        if asgraph:
            return _H1toTGraph(yoda_obj, widthscale, usefocus, includeOverflows, includeHiddenBins)
        else:
            return _H1toTH1D(yoda_obj, widthscale)
    elif isinstance(yoda_obj, yoda.BinnedProfile1D):
        if asgraph:
            return _P1toTGraph(yoda_obj, widthscale, usefocus, includeOverflows, includeHiddenBins)
        else:
            return _P1toTProfile(yoda_obj)
    elif isinstance(yoda_obj, yoda.BinnedEstimate1D):
        if asgraph:
            return _E1toTGraph(yoda_obj, includeOverflows, includeHiddenBins)
        else:
            return _E1toTH1D(yoda_obj)
    elif isinstance(yoda_obj, yoda.BinnedHisto2D):
        return _H2toTH2D(yoda_obj, widthscale)
    elif isinstance(yoda_obj, yoda.BinnedEstimate2D):
        return _E2toTH2D(yoda_obj)
    elif isinstance(yoda_obj, yoda.Scatter2D):
        return _S2toTGraph(yoda_obj)
    else:
        if yoda_obj.dim() == 1:
            return _S1toTGraph(yoda_obj.mkScatter(yoda_obj.path()))
        elif yoda_obj.dim() == 2:
            return _S2toTGraph(yoda_obj.mkScatter(yoda_obj.path()))
        else:
            raise Exception(f"Unknown/unsupported type for conversion: {type(yoda_obj)}")
