# 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.



"""

import os
import exceptions
import numpy as np
import h5py as h5
from types import NoneType
import re

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


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


class FileError(exceptions.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(exceptions.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(exceptions.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(exceptions.Exception):
    """ An exception raised for non-implemented functionality"""

    def __init__(self):
        return

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


class H5GroupWrapper(h5.Group):
    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):
    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.  """
    # 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


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 = {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, **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.
    """

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

        _dump(py_obj, h5f, **kwargs)
        h5f.close()
    except NoMatchError:
        fname = h5f.filename
        h5f.close()
        try:
            os.remove(fname)
        except:
            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.  """
    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,
        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)
    #print call_id, create_dataset

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
    """
    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"""
    d = h_group.create_dataset('data_%i' % call_id, data=py_obj)
    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"""
    d = h_group.create_dataset('data_%i' % call_id, data=py_obj, dtype=type(py_obj))
    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
    """
    h_dictgroup = h_group.create_group('data_%i' % call_id)
    h_dictgroup.attrs["type"] = ['dict']
    for key, py_subobj in py_obj.items():
        #print key, py_subobj, h_group
        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"""
    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)
        d.attrs["mask"] = py_obj.mask
        d.attrs["type"] = ['ndarray_masked']
        #m.attrs["type"] = ['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"""

    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)
        ll = len(py_obj)
        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 """
    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 """
    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 ##
#############

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

    Args:
    file: 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)
    """

    try:
        h5f = file_opener(file)
        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, h5f)
            return py_container[0][0]
        except AssertionError:
            import hickle_legacy
            hickle_legacy.load(file, safe)

    finally:
        if 'h5f' in locals():
            h5f.close()


def load_dataset(h_node):
    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':
        try:
            mask = h_node.attrs["mask"][:]
        except IndexError:
            mask = h_node.attrs["mask"]
        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)

        print subtype, tcast, data

        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
    """
    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 """

    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

class PyContainer(list):
    def __init__(self):
        super(PyContainer, self).__init__()
        self.container_type = None
        self.name = None

    def convert(self):
        #print self.container_type

        if self.container_type == "<type 'list'>":
            #print self.name
            return list(self)
        if self.container_type == "<type 'tuple'>":
            #print self.name
            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


