import h5py


def loadDataset(filename):
    hid = h5py.File(filename, "r")
    icount = 0
    snap = loadSnap(hid)
    data={}
    act= {}
    maxID = 0
    for key in hid.keys():
        if 'scan' in key:
            ID = int(key.split('_')[1])
            if ID > maxID:
                maxID = ID
            print('Reading scan_%d' % ID)
            data[ID]=loadData(hid,ID)
            act[ID]=loadActuator(hid,ID)
            icount +=1
    hid.close()
    if icount == 0:
        return None, None, None
    elif icount == 1:
        return data[0], act[0],snap
    return [data[j] for j in range(maxID+1)],[act[j] for j in range(maxID+1)],snap
 
def loadSnap(hid):
    snap={}
    if not 'experiment' in hid.keys():
        return None
    for key1 in hid['experiment'].keys():
        if isinstance(hid['experiment'][key1],h5py.Group):
            for key2 in hid['experiment'][key1].keys():
                val = hid['experiment'][key1][key2][()]
                snap[key1+':'+key2]=val
        else:
            snap[key1]=hid['experiment'][key1][()]
    return snap

def loadData(hid,scanrun=1):
    run='scan_%d' % scanrun
    data = {} 
    for key1 in hid[run]['data'].keys():
        if isinstance(hid[run]['data'][key1],h5py.Group):
            for key2 in hid[run]['data'][key1].keys():
                val = hid[run]['data'][key1][key2][()]
                data[key1+':'+key2]=val
        else:
            data[key1]=hid[run]['data'][key1][()]
    return data

def loadActuator(hid,scanrun=1):
    run='scan_%d' % scanrun
    data = {} 
    if 'actuators' in hid[run]['method'].keys():
        for key1 in hid[run]['method']['actuators'].keys():
            if isinstance(hid[run]['method']['actuators'][key1],h5py.Group):
                for key2 in hid[run]['method']['actuators'][key1].keys():
                    val = hid[run]['method']['actuators'][key1][key2][()]
                    data[key1+':'+key2]={'val':val}
            else:
                data[key1]=hid[run]['method']['actuators'][key1][()]
    return data