import os
import sys
import tempfile
import unittest

import mock
import numpy

from chainer import cuda
from chainer import link
from chainer import links
from chainer import optimizers
from chainer.serializers import hdf5
from chainer import testing
from chainer.testing import attr


if hdf5._available:
    import h5py


@unittest.skipUnless(hdf5._available, 'h5py is not available')
class TestHDF5Serializer(unittest.TestCase):

    def setUp(self):
        fd, path = tempfile.mkstemp()
        os.close(fd)
        self.temp_file_path = path
        self.hdf5file = h5py.File(path, 'w')
        self.serializer = hdf5.HDF5Serializer(self.hdf5file, compression=3)

        self.data = numpy.random.uniform(-1, 1, (2, 3)).astype(numpy.float32)

    def tearDown(self):
        if hasattr(self, 'hdf5file'):
            self.hdf5file.close()
        if hasattr(self, 'temp_file_path'):
            os.remove(self.temp_file_path)

    def test_get_item(self):
        child = self.serializer['x']
        self.assertIsInstance(child, hdf5.HDF5Serializer)
        self.assertEqual(child.group.name, '/x')
        self.assertEqual(child.compression, 3)

    def check_serialize(self, data):
        ret = self.serializer('w', data)
        dset = self.hdf5file['w']

        self.assertIsInstance(dset, h5py.Dataset)
        self.assertEqual(dset.shape, data.shape)
        self.assertEqual(dset.size, data.size)
        self.assertEqual(dset.dtype, data.dtype)
        read = numpy.empty((2, 3), dtype=numpy.float32)
        dset.read_direct(read)
        numpy.testing.assert_array_equal(read, cuda.to_cpu(data))

        self.assertEqual(dset.compression_opts, 3)

        self.assertIs(ret, data)

    def test_serialize_cpu(self):
        self.check_serialize(self.data)

    @attr.gpu
    def test_serialize_gpu(self):
        self.check_serialize(cuda.to_gpu(self.data))

    def test_serialize_scalar(self):
        ret = self.serializer('x', 10)
        dset = self.hdf5file['x']

        self.assertIsInstance(dset, h5py.Dataset)
        self.assertEqual(dset.shape, ())
        self.assertEqual(dset.size, 1)
        self.assertEqual(dset.dtype, int)
        read = numpy.empty((), dtype=numpy.int32)