import os

import numpy

from chainer import cuda
from chainer import serializer


try:
    import h5py
except ImportError:
    if not os.environ.get('READTHEDOCS', None) == 'True':
        raise


class HDF5Serializer(serializer.Serializer):

    """Serializer for HDF5 format.

    This is the standard serializer in Chainer. The chain hierarchy is simply
    mapped to HDF5 hierarchical groups.

    Args:
        group (h5py.Group): The group that this serializer represents.
        compression (int): Gzip compression level.

    """
    def __init__(self, group, compression=4):
        self.group = group
        self.compression = compression

    def __getitem__(self, key):
        return HDF5Serializer(self.group.require_group(key), self.compression)

    def __call__(self, key, value):
        ret = value
        if isinstance(value, cuda.ndarray):
            value = cuda.to_cpu(value)
        arr = numpy.asarray(value)
        compression = None if arr.size <= 1 else self.compression
        self.group.create_dataset(key, data=arr, compression=compression)
        return ret


def save_hdf5(filename, obj, compression=4):
    """Saves an object to the file in HDF5 format.

    This is a short-cut function to save only one object into an HDF5 file. If
    you want to save multiple objects to one HDF5 file, use
    :class:`HDF5Serializer` directly by passing appropriate :class:`h5py.Group`
    objects.

    Args:
        filename (str): Target file name.
        obj: Object to be serialized. It must support serialization protocol.
        compression (int): Gzip compression level.

    """
    f = h5py.File(filename, 'w')
    s = HDF5Serializer(f, compression=compression)
    s.save(obj)


class HDF5Deserializer(serializer.Deserializer):

    """Deserializer for HDF5 format.

    This is the standard deserializer in Chainer. This deserializer can be used
    to read an object serialized by :class:`HDF5Serializer`.

    Args:
        group (h5py.Group): The group that the deserialization starts from.

    """
    def __init__(self, group):
        self.group = group

    def __getitem__(self, key):
        return HDF5Deserializer(self.group[key])

    def __call__(self, key, value):
        dataset = self.group[key]
        if isinstance(value, numpy.ndarray):
            dataset.read_direct(value)
        elif isinstance(value, cuda.ndarray):
            value.set(numpy.asarray(dataset))
        else:
            value = type(value)(numpy.asarray(dataset))
        return value


def load_hdf5(filename, obj):
    """Loads an object from the file in HDF5 format.

    This is a short-cut function to load from an HDF5 file that contains only
    one object. If you want to load multiple objects from one HDF5 file, use
    :class:`HDF5Deserializer` directly by passing appropriate
    :class:`h5py.Group` objects.

    Args:
        filename (str): Name of the file to be loaded.
        obj: Object to be deserialized. It must support serialization protocol.

    """
    f = h5py.File(filename, 'r')
    d = HDF5Deserializer(f)
    d.load(obj)
