import json
import os
import ROOT as r
import copy
import math

def nameToDict(name):
    f = open('SFs/'+name,'r')
    dct = json.load(f)
    f.close()
    return dct 

def elemToDicts(elem):
    print 'stats from', elem[0][0]
    stat = nameToDict(elem[0][0])
    syst = {} 
    for id in elem[1]:
        print 'systs from', id
        syst.update( nameToDict(id))

    for nam in stat:
        if nam not in syst: raise RuntimeError('Name %s not in syst'%nam)
    for nam in syst:
        if nam not in stat: raise RuntimeError('Name %s not in stat'%nam)

    return stat,syst


def getTriplet(stat, syst):
    nom = stat
    outStat = stat.Clone(nom.GetName() + '_stat')
    outSyst = syst.Clone(nom.GetName() + '_syst')    



    for x in range(1,nom.GetXaxis().GetNbins()+1):
        for y in range(1,nom.GetYaxis().GetNbins()+1):
            bin = nom.GetBin(x,y)
            errStat = outStat.GetBinError(bin)
            errSyst = outSyst.GetBinError(bin)
            nom.SetBinError(bin, math.sqrt(errSyst**2+ errStat**2))
    return (nom, outStat, outSyst)


def mergeElem(elem):
    stat,syst = elemToDicts(elem)
    for id in stat:
        for binning in stat[id]:
            for var1 in stat[id][binning]:
                for var2 in stat[id][binning][var1]:
                    #print id, binning, var1, var2
                    if not stat[id][binning][var1][var2]['value']: raise RuntimeError('calm down boy')
                    if syst[id][binning][var1][var2]['value'] != 'null' and abs(float(syst[id][binning][var1][var2]['value'])-float(stat[id][binning][var1][var2]['value'])) > 5e-3:
                        #print syst[id][binning][var1][var2]['value'], stat[id][binning][var1][var2]['value']
                        raise RuntimeError('damn it')
                    tstat = stat[id][binning][var1][var2]['error']
                    tsyst = syst[id][binning][var1][var2]['error']
                    del stat[id][binning][var1][var2]['error']
                    stat[id][binning][var1][var2]['stat'] = tstat
                    stat[id][binning][var1][var2]['syst'] = tsyst

    kk= json.dumps(stat, indent=4, sort_keys=True)
    outfile = open('out/' + elem[0][0], 'w')
    outfile.write(kk)

def mergeRoots(elem):
    # first merge systematics 
    os.system('hadd -f SFs/kk/{name} {lst}'.format(name=elem[0][0].replace('.json','.root'),
                                                      lst =' '.join(['SFs/'+x.replace('.json','.root') for x in elem[1]])))
    tsyst = r.TFile.Open('SFs/kk/%s'%elem[0][0].replace('.json','.root'))
    tstat = r.TFile.Open('SFs/%s'%elem[0][0].replace('.json','.root'))


    hStat = []; hSyst=[]

    for hist in tstat.GetListOfKeys():
        if not hist.ReadObj().InheritsFrom("TH2D"): raise 
        hStat.append(copy.deepcopy(hist.ReadObj()))
        hStat[-1].GetBinError(2,2)
    for hst in hStat:
        hSyst.append(copy.deepcopy(tsyst.Get(hst.GetName())))
        hSyst[-1].GetBinError(2,2)
    tsyst.Close()
    tstat.Close()

    triplets = [] 
    for i in range(len(hStat)):
        print hStat[i].GetBinError(2,2), hSyst[i].GetBinError(2,2)
        triplets.append(getTriplet( hStat[i], hSyst[i]))

    tfile = r.TFile.Open('out/%s'%elem[0][0].replace('.json','.root'),'recreate')
    for tri in triplets:
        tri[0].Write(); tri[1].Write(); tri[2].Write()
    tfile.Close()


jsonList = [
    [ ['RunBCDEF_SF_ID_JPsi.json'],
      ['RunBCDEF_JPsi_sf_SYS_Loose.json',
       'RunBCDEF_JPsi_sf_SYS_Medium.json',
       'RunBCDEF_JPsi_sf_SYS_Soft.json',
       'RunBCDEF_JPsi_sf_SYS_Tight.json']],
    [ ['RunBCDEF_SF_ID.json'],
      ['RunBCDEF_sf_SYS_NUM_HighPtID_DEN_genTracks.json',
       'RunBCDEF_sf_SYS_NUM_LooseID_DEN_genTracks.json',
       'RunBCDEF_sf_SYS_NUM_MediumID_DEN_genTracks.json',
       'RunBCDEF_sf_SYS_NUM_MediumPromptID_DEN_genTracks.json',
       'RunBCDEF_sf_SYS_NUM_SoftID_DEN_genTracks.json',
       'RunBCDEF_sf_SYS_NUM_TightID_DEN_genTracks.json',
       'RunBCDEF_sf_SYS_NUM_TrkHighPtID_DEN_genTracks.json']],
    [ ['RunBCDEF_SF_ISO.json'],
      ['RunBCDEF_sf_SYS_NUM_TightRelIso_DEN_MediumID.json',
       'RunBCDEF_sf_SYS_NUM_TightRelIso_DEN_TightIDandIPCut.json',
       'RunBCDEF_sf_SYS_NUM_TightRelTkIso_DEN_HighPtIDandIPCut.json',
       'RunBCDEF_sf_SYS_NUM_TightRelTkIso_DEN_TrkHighPtID.json',
       'RunBCDEF_sf_SYS_NUM_LooseRelIso_DEN_LooseID.json',
       'RunBCDEF_sf_SYS_NUM_LooseRelIso_DEN_MediumID.json',
       'RunBCDEF_sf_SYS_NUM_LooseRelIso_DEN_TightIDandIPCut.json',
       'RunBCDEF_sf_SYS_NUM_LooseRelTkIso_DEN_HighPtIDandIPCut.json',
       'RunBCDEF_sf_SYS_NUM_LooseRelTkIso_DEN_TrkHighPtID.json',
   ]]
]
              
#for typ in jsonList:
#    mergeElem(typ)
              
for typ in jsonList:
    mergeRoots(typ)