# encoding: utf-8
"""
# hickle.py

Created by Danny Price 2016-02-03.

Hickle is a HDF5 based clone of Pickle. Instead of serializing to a pickle
file, Hickle dumps to a HDF5 file. It is designed to be as similar to pickle in
usage as possible, providing a load() and dump() function.

## Notes

Hickle has two main advantages over Pickle:
1) LARGE PICKLE HANDLING. Unpickling a large pickle is slow, as the Unpickler
reads the entire pickle thing and loads it into memory. In comparison, HDF5
files are designed for large datasets. Things are only loaded when accessed.

2) CROSS PLATFORM SUPPORT. Attempting to unpickle a pickle pickled on Windows
on Linux and vice versa is likely to fail with errors like "Insecure string
pickle". HDF5 files will load fine, as long as both machines have
h5py installed.

"""

import os
import numpy as np
import h5py as h5
import re

try:
    from exceptions import Exception
    from types import NoneType
except ImportError:
    pass        # above imports will fail in python3

import warnings
__version__ = "2.0.4"
__author__ = "Danny Price"


##################
# Error handling #
##################

class FileError(Exception):
    """ An exception raised if the file is fishy """
    def __init__(self):
        return

    def __str__(self):
        return ("Cannot open file. Please pass either a filename "
                "string, a file object, or a h5py.File")


class ClosedFileError(Exception):
    """ An exception raised if the file is fishy """
    def __init__(self):
        return

    def __str__(self):
        return ("HDF5 file has been closed. Please pass either "
                "a filename string, a file object, or an open h5py.File")


class NoMatchError(Exception):
    """ An exception raised if the object type is not understood (or
    supported)"""
    def __init__(self):
        return

    def __str__(self):
        return ("Error: this type of python object cannot be converted into a "
                "hickle.")


class ToDoError(Exception):
    """ An exception raised for non-implemented functionality"""
    def __init__(self):
        return

    def __str__(self):
        return "Error: this functionality hasn't been implemented yet."


######################
# H5PY file wrappers #
######################

class H5GroupWrapper(h5.Group):
    """ Group wrapper that provides a track_times kwarg.

    track_times is a boolean flag that can be set to False, so that two
    files created at different times will have identical MD5 hashes.
    """
    def create_dataset(self, *args, **kwargs):
        kwargs['track_times'] = getattr(self, 'track_times', True)
        return super(H5GroupWrapper, self).create_dataset(*args, **kwargs)

    def create_group(self, *args, **kwargs):
        group = super(H5GroupWrapper, self).create_group(*args, **kwargs)
        group.__class__ = H5GroupWrapper
        group.track_times = getattr(self, 'track_times', True)
        return group


class H5FileWrapper(h5.File):
    """ Wrapper for h5py File that provides a track_times kwarg.

    track_times is a boolean flag that can be set to False, so that two
    files created at different times will have identical MD5 hashes.
    """
    def create_dataset(self, *args, **kwargs):
        kwargs['track_times'] = getattr(self, 'track_times', True)
        return super(H5FileWrapper, self).create_dataset(*args, **kwargs)

    def create_group(self, *args, **kwargs):
        group = super(H5FileWrapper, self).create_group(*args, **kwargs)
        group.__class__ = H5GroupWrapper
        group.track_times = getattr(self, 'track_times', True)
        return group


def file_opener(f, mode='r', track_times=True):
    """ A file opener helper function with some error handling.  This can open
    files through a file object, a h5py file, or just the filename.

    Args:
        f (file, h5py.File, or string): File-identifier, e.g. filename or file object.
        mode (str): File open mode. Only required if opening by filename string.
        track_times (bool): Track time in HDF5; turn off if you want hickling at
                 different times to produce identical files (e.g. for MD5 hash check).

    """
    # Were we handed a file object or just a file name string?
    if isinstance(f, file):
        filename, mode = f.name, f.mode
        f.close()
        h5f = h5.File(filename, mode)
    elif isinstance(f, str) or isinstance(f, unicode):
        filename = f
        h5f = h5.File(filename, mode)
    elif isinstance(f, H5FileWrapper) or isinstance(f, h5._hl.files.File):
        try:
            filename = f.filename
        except ValueError:
            raise ClosedFileError()
        h5f = f
    else:
        print(type(f))
        raise FileError

    h5f.__class__ = H5FileWrapper
    h5f.track_times = track_times
    return h5f


###########
# DUMPERS #
###########

def check_is_iterable(py_obj):
    """ Check whether a python object is iterable.

    Note: this treats unicode and string as NON ITERABLE

    Args:
        py_obj: python object to test

    Returns:
        iter_ok (bool): True if item is iterable, False is item is not
    """
    if type(py_obj) in (str, unicode):
        return False
    try:
        iter(py_obj)
        return True
    except TypeError:
        return False


def check_iterable_item_type(iter_obj):
    """ Check if all items within an iterable are the same type.

    Args:
        iter_obj: iterable object

    Returns:
        iter_type: type of item contained within the iterable. If
                   the iterable has many types, a boolean False is returned instead.

    References:
    http://stackoverflow.com/questions/13252333/python-check-if-all-elements-of-a-list-are-the-same-type
    """
    iseq = iter(iter_obj)
    first_type = type(next(iseq))
    return first_type if all((type(x) is first_type) for x in iseq) else False


def check_is_numpy_array(py_obj):
    """ Check if a python object is a numpy array (masked or regular)

    Args:
        py_obj: python object to check whether it is a numpy array

    Returns
        is_numpy (bool): Returns True if it is a numpy array, else False if it isn't
    """

    is_numpy = type(py_obj) in (type(np.array([1])), type(np.ma.array([1])))

    return is_numpy


def _dump(py_obj, h_group, call_id=0, **kwargs):
    """ Dump a python object to a group within a HDF5 file.

    This function is called recursively by the main dump() function.

    Args:
        py_obj: python object to dump.
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.
    """

    dumpable_dtypes = set([bool, int, float, long, complex, str, unicode])

    # Firstly, check if item is a numpy array. If so, just dump it.
    if check_is_numpy_array(py_obj):
        create_hkl_dataset(py_obj, h_group, call_id, **kwargs)

    # next, check if item is iterable
    elif check_is_iterable(py_obj):
        item_type = check_iterable_item_type(py_obj)

        # item_type == False implies multiple types. Create a dataset
        if item_type is False:
            h_subgroup = create_hkl_group(py_obj, h_group, call_id)
            for ii, py_subobj in enumerate(py_obj):
                _dump(py_subobj, h_subgroup, call_id=ii, **kwargs)

        # otherwise, subitems have same type. Check if subtype is an iterable
        # (e.g. list of lists), or not (e.g. list of ints, which should be treated
        # as a single dataset).
        else:
            if item_type in dumpable_dtypes:
                create_hkl_dataset(py_obj, h_group, call_id, **kwargs)
            else:
                h_subgroup = create_hkl_group(py_obj, h_group, call_id)
                for ii, py_subobj in enumerate(py_obj):
                    #print py_subobj, h_subgroup, ii
                    _dump(py_subobj, h_subgroup, call_id=ii, **kwargs)

    # item is not iterable, so create a dataset for it
    else:
        create_hkl_dataset(py_obj, h_group, call_id, **kwargs)


def dump(py_obj, file_obj, mode='w', track_times=True, path='/', **kwargs):
    """ Write a pickled representation of obj to the open file object file.

    Args:
    obj (object): python object o store in a Hickle
    file: file object, filename string, or h5py.File object
            file in which to store the object. A h5py.File or a filename is also
            acceptable.
    mode (str): optional argument, 'r' (read only), 'w' (write) or 'a' (append).
            Ignored if file is a file object.
    compression (str): optional argument. Applies compression to dataset. Options: None, gzip,
            lzf (+ szip, if installed)
    track_times (bool): optional argument. If set to False, repeated hickling will produce
            identical files.
    path (str): path within hdf5 file to save data to. Defaults to root /
    """

    try:
        # Open the file
        h5f = file_opener(file_obj, mode, track_times)
        h5f.attrs["CLASS"] = 'hickle'
        h5f.attrs["VERSION"] = 2
        h5f.attrs["type"] = ['hickle']

        h_root_group = h5f.get(path)

        if h_root_group is None:
            h_root_group = h5f.create_group(path)
            h_root_group.attrs["type"] = ['hickle']

        _dump(py_obj, h_root_group, **kwargs)
        h5f.close()
    except NoMatchError:
        fname = h5f.filename
        h5f.close()
        try:
            os.remove(fname)
        except OSError:
            warnings.warn("Dump failed. Could not remove %s" % fname)
        finally:
            raise NoMatchError


def create_dataset_lookup(py_obj):
    """ What type of object are we trying to pickle?  This is a python
    dictionary based equivalent of a case statement.  It returns the correct
    helper function for a given data type.

    Args:
        py_obj: python object to look-up what function to use to dump to disk

    Returns:
        match: function that should be used to dump data to a new dataset
    """
    t = type(py_obj)

    types = {
        dict: create_dict_dataset,
        list: create_listlike_dataset,
        tuple: create_listlike_dataset,
        set: create_listlike_dataset,
        str: create_stringlike_dataset,
        unicode: create_stringlike_dataset,
        int: create_python_dtype_dataset,
        float: create_python_dtype_dataset,
        long: create_python_dtype_dataset,
        bool: create_python_dtype_dataset,
        complex: create_python_dtype_dataset,
        NoneType: create_none_dataset,
        np.ndarray: create_np_array_dataset,
        np.ma.core.MaskedArray: create_np_array_dataset,
        np.float16: create_np_dtype_dataset,
        np.float32: create_np_dtype_dataset,
        np.float64: create_np_dtype_dataset,
        np.int8: create_np_dtype_dataset,
        np.int16: create_np_dtype_dataset,
        np.int32: create_np_dtype_dataset,
        np.int64: create_np_dtype_dataset,
        np.uint8: create_np_dtype_dataset,
        np.uint16: create_np_dtype_dataset,
        np.uint32: create_np_dtype_dataset,
        np.uint64: create_np_dtype_dataset,
        np.complex64: create_np_dtype_dataset,
        np.complex128: create_np_dtype_dataset
    }

    match = types.get(t, no_match)
    return match


def create_hkl_dataset(py_obj, h_group, call_id=0, **kwargs):
    """ Create a dataset within the hickle HDF5 file

    Args:
        py_obj: python object to dump.
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.

    """
    #lookup dataset creator type based on python object type
    create_dataset = create_dataset_lookup(py_obj)

    # do the creation
    create_dataset(py_obj, h_group, call_id, **kwargs)


def create_hkl_group(py_obj, h_group, call_id=0):
    """ Create a new group within the hickle file

    Args:
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.

    """
    h_subgroup = h_group.create_group('data_%i' % call_id)
    h_subgroup.attrs["type"] = [str(type(py_obj))]
    return h_subgroup


def create_listlike_dataset(py_obj, h_group, call_id=0, **kwargs):
    """ Dumper for list, set, tuple

    Args:
        py_obj: python object to dump; should be list-like
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.
    """
    dtype = str(type(py_obj))
    obj = list(py_obj)
    d = h_group.create_dataset('data_%i' % call_id, data=obj, **kwargs)
    d.attrs["type"] = [dtype]


def create_np_dtype_dataset(py_obj, h_group, call_id=0, **kwargs):
    """ dumps an np dtype object to h5py file

    Args:
        py_obj: python object to dump; should be a numpy scalar, e.g. np.float16(1)
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.
    """
    d = h_group.create_dataset('data_%i' % call_id, data=py_obj, **kwargs)
    d.attrs["type"] = ['np_dtype']
    d.attrs["np_dtype"] = str(d.dtype)


def create_python_dtype_dataset(py_obj, h_group, call_id=0, **kwargs):
    """ dumps a python dtype object to h5py file

    Args:
        py_obj: python object to dump; should be a python type (int, float, bool etc)
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.
    """
    d = h_group.create_dataset('data_%i' % call_id, data=py_obj,
                               dtype=type(py_obj), **kwargs)
    d.attrs["type"] = ['python_dtype']
    d.attrs['python_subdtype'] = str(type(py_obj))


def create_dict_dataset(py_obj, h_group, call_id=0, **kwargs):
    """ Creates a data group for each key in dictionary

    Args:
        py_obj: python object to dump; should be dictionary
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.
    """
    h_dictgroup = h_group.create_group('data_%i' % call_id)
    h_dictgroup.attrs["type"] = ['dict']
    for key, py_subobj in py_obj.items():
        h_subgroup = h_dictgroup.create_group(key)
        h_subgroup.attrs["type"] = ['dict_item']
        _dump(py_subobj, h_subgroup, call_id=0, **kwargs)


def create_np_array_dataset(py_obj, h_group, call_id=0, **kwargs):
    """ dumps an ndarray object to h5py file

    Args:
        py_obj: python object to dump; should be a numpy array or np.ma.array (masked)
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.
    """
    if isinstance(py_obj, type(np.ma.array([1]))):
        d = h_group.create_dataset('data_%i' % call_id, data=py_obj, **kwargs)
        #m = h_group.create_dataset('mask_%i' % call_id, data=py_obj.mask, **kwargs)
        m = h_group.create_dataset('data_%i_mask' % call_id, data=py_obj.mask, **kwargs)
        d.attrs["type"] = ['ndarray_masked_data']
        m.attrs["type"] = ['ndarray_masked_mask']
    else:
        d = h_group.create_dataset('data_%i' % call_id, data=py_obj, **kwargs)
        d.attrs["type"] = ['ndarray']


def create_stringlike_dataset(py_obj, h_group, call_id=0, **kwargs):
    """ dumps a list object to h5py file

    Args:
        py_obj: python object to dump; should be string-like (unicode or string)
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.
    """
    if isinstance(py_obj, str):
        d = h_group.create_dataset('data_%i' % call_id, data=[py_obj], **kwargs)
        d.attrs["type"] = ['string']
    else:
        dt = h5.special_dtype(vlen=unicode)
        dset = h_group.create_dataset('data_%i' % call_id, shape=(1, ), dtype=dt, **kwargs)
        dset[0] = py_obj
        dset.attrs['type'] = ['unicode']


def create_none_dataset(py_obj, h_group, call_id=0, **kwargs):
    """ Dump None type to file

    Args:
        py_obj: python object to dump; must be None object
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.
    """
    d = h_group.create_dataset('data_%i' % call_id, data=[0], **kwargs)
    d.attrs["type"] = ['none']


def no_match(py_obj, h_group, call_id=0, **kwargs):
    """ If no match is made, raise an exception

    Args:
        py_obj: python object to dump; default if item is not matched.
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.
    """
    import cPickle

    pickled_obj = cPickle.dumps(py_obj)
    d = h_group.create_dataset('data_%i' % call_id, data=[pickled_obj])
    d.attrs["type"] = ['pickle']

    warnings.warn("%s type not understood, data have been "
                  "serialized" % type(py_obj))


#############
## LOADERS ##
#############

class PyContainer(list):
    """ A group-like object into which to load datasets.

    In order to build up a tree-like structure, we need to be able
    to load datasets into a container with an append() method.
    Python tuples and sets do not allow this. This class provides
    a list-like object that be converted into a list, tuple, set or dict.
    """
    def __init__(self):
        super(PyContainer, self).__init__()
        self.container_type = None
        self.name = None

    def convert(self):
        """ Convert from PyContainer to python core data type.

        Returns: self, either as a list, tuple, set or dict
        """
        if self.container_type == "<type 'list'>":
            return list(self)
        if self.container_type == "<type 'tuple'>":
            return tuple(self)
        if self.container_type == "<type 'set'>":
            return set(self)
        if self.container_type == "dict":
            keys = [str(item.name.split('/')[-1]) for item in self]
            items = [item[0] for item in self]
            return dict(zip(keys, items))
        else:
            return self


def load(fileobj, path='/', safe=True):
    """ Load a hickle file and reconstruct a python object

    Args:
        fileobj: file object, h5py.File, or filename string
            safe (bool): Disable automatic depickling of arbitrary python objects.
            DO NOT set this to False unless the file is from a trusted source.
            (see http://www.cs.jhu.edu/~s/musings/pickle.html for an explanation)

        path (str): path within hdf5 file to save data to. Defaults to root /
    """

    try:
        h5f = file_opener(fileobj)

        h_root_group = h5f.get(path)
        try:
            assert 'CLASS' in h5f.attrs.keys()
            assert 'VERSION' in h5f.attrs.keys()
            py_container = PyContainer()
            py_container.container_type = 'hickle'
            py_container = _load(py_container, h_root_group)
            return py_container[0][0]
        except AssertionError:
            import hickle_legacy
            return hickle_legacy.load(fileobj, safe)
    finally:
        if 'h5f' in locals():
            h5f.close()


def load_dataset(h_node):
    """ Load a dataset, converting into its correct python type

    Args:
        h_node (h5py dataset): h5py dataset object to read

    Returns:
        data: reconstructed python object from loaded data
    """
    py_type = h_node.attrs["type"][0]

    if h_node.shape == ():
        data = h_node.value
    else:
        data  = h_node[:]

    if py_type == "<type 'list'>":
        #print self.name
        return list(data)
    elif py_type == "<type 'tuple'>":
        return tuple(data)
    elif py_type == "<type 'set'>":
        return set(data)
    elif py_type == "np_dtype":
        subtype = h_node.attrs["np_dtype"]
        data = np.array(data, dtype=subtype)
        return data
    elif py_type == 'ndarray':
        return np.array(data)
    elif py_type == 'ndarray_masked_data':
        try:
            mask_path = h_node.name + "_mask"
            h_root = h_node.parent
            mask = h_root.get(mask_path)[:]
        except IndexError:
            mask = h_root.get(mask_path)
        except ValueError:
            mask = h_root.get(mask_path)
        data = np.ma.array(data, mask=mask)
        return data
    elif py_type == 'python_dtype':
        subtype = h_node.attrs["python_subdtype"]
        type_dict = {
            "<type 'int'>": int,
            "<type 'float'>": float,
            "<type 'long'>": long,
            "<type 'bool'>": bool,
            "<type 'complex'>": complex
        }
        tcast = type_dict.get(subtype)
        return tcast(data)
    elif py_type == 'string':
        return str(data[0])
    elif py_type == 'unicode':
        return unicode(data[0])
    elif py_type == 'none':
        return None
    else:
        print(h_node.name, py_type, h_node.attrs.keys())
        return data


def sort_keys(key_list):
    """ Take a list of strings and sort it by integer value within string

    Args:
        key_list (list): List of keys

    Returns:
        key_list_sorted (list): List of keys, sorted by integer
    """
    to_int = lambda x: int(re.search('\d+', x).group(0))
    keys_by_int = sorted([(to_int(key), key) for key in key_list])
    return [ii[1] for ii in keys_by_int]


def _load(py_container, h_group):
    """ Load a hickle file

    Recursive funnction to load hdf5 data into a PyContainer()

    Args:
        py_container (PyContainer): Python container to load data into
        h_group (h5 group or dataset): h5py object, group or dataset, to spider
                and load all datasets.
    """

    group_dtype   = h5._hl.group.Group
    dataset_dtype = h5._hl.dataset.Dataset

    #either a file, group, or dataset
    if isinstance(h_group, H5FileWrapper) or isinstance(h_group, group_dtype):
        py_subcontainer = PyContainer()
        py_subcontainer.container_type = h_group.attrs['type'][0]
        py_subcontainer.name = h_group.name

        if py_subcontainer.container_type != 'dict':
            h_keys = sort_keys(h_group.keys())
        else:
            h_keys = h_group.keys()

        for h_name in h_keys:
            h_node = h_group[h_name]
            py_subcontainer = _load(py_subcontainer, h_node)

        sub_data = py_subcontainer.convert()
        py_container.append(sub_data)

    else:
        # must be a dataset
        subdata = load_dataset(h_group)
        py_container.append(subdata)

    #print h_group.name, py_container
    return py_container
