import atexit
import shutil
import tempfile
from multiprocessing import Pool as ProcessPool
from multiprocessing import cpu_count
from multiprocessing.pool import ThreadPool
from tempfile import mkdtemp

import numpy as np
from numpy.testing import assert_array_equal

from zarr.attrs import Attributes
from zarr.core import Array
from zarr.hierarchy import Group
from zarr.storage import (DirectoryStore, KVStore, atexit_rmtree, init_array,
                          init_group, meta_root)
from zarr.sync import ProcessSynchronizer, ThreadSynchronizer
# zarr_version fixture must be imported although not used directly here
from zarr.tests.test_attrs import TestAttributes, zarr_version  # noqa
from zarr.tests.test_core import TestArray
from zarr.tests.test_hierarchy import TestGroup


class TestAttributesWithThreadSynchronizer(TestAttributes):

    def init_attributes(self, store, read_only=False, cache=True, zarr_version=zarr_version):
        key = '.zattrs' if zarr_version == 2 else meta_root + 'attrs'
        synchronizer = ThreadSynchronizer()
        return Attributes(store, synchronizer=synchronizer, key=key,
                          read_only=read_only, cache=cache)


class TestAttributesProcessSynchronizer(TestAttributes):

    def init_attributes(self, store, read_only=False, cache=True, zarr_version=zarr_version):
        key = '.zattrs' if zarr_version == 2 else meta_root + 'attrs'
        sync_path = mkdtemp()
        atexit.register(shutil.rmtree, sync_path)
        synchronizer = ProcessSynchronizer(sync_path)
        return Attributes(store, synchronizer=synchronizer, key=key,
                          read_only=read_only, cache=cache)


def _append(arg):
    z, i = arg
    import numpy
    x = numpy.empty(1000, dtype='i4')
    x[:] = i
    shape = z.append(x)
    return shape


def _set_arange(arg):
    z, i = arg
    import numpy
    x = numpy.arange(i*1000, (i*1000)+1000, 1)
    z[i*1000:(i*1000)+1000] = x
    return i


class MixinArraySyncTests:

    def test_parallel_setitem(self):
        n = 100

        # setup
        arr = self.create_array(shape=n * 1000, chunks=999, dtype='i4')
        arr[:] = 0
        pool = self.create_pool()

        # parallel setitem
        results = pool.map(_set_arange, zip([arr] * n, range(n)), chunksize=1)
        results = sorted(results)

        assert list(range(n)) == results
        assert_array_equal(np.arange(n * 1000), arr[:])

        pool.terminate()

    def test_parallel_append(self):
        n = 100

        # setup
        arr = self.create_array(shape=1000, chunks=999, dtype='i4')
        arr[:] = 0
        pool = self.create_pool()

        # parallel append
        results = pool.map(_append, zip([arr] * n, range(n)), chunksize=1)
        results = sorted(results)

        assert [((i+2)*1000,) for i in range(n)] == results
        assert ((n+1)*1000,) == arr.shape

        pool.terminate()


class TestArrayWithThreadSynchronizer(TestArray, MixinArraySyncTests):

    def create_array(self, read_only=False, **kwargs):
        store = KVStore(dict())
        cache_metadata = kwargs.pop('cache_metadata', True)
        cache_attrs = kwargs.pop('cache_attrs', True)
        write_empty_chunks = kwargs.pop('write_empty_chunks', True)
        init_array(store, **kwargs)
        return Array(store, synchronizer=ThreadSynchronizer(),
                     read_only=read_only, cache_metadata=cache_metadata,
                     cache_attrs=cache_attrs, write_empty_chunks=write_empty_chunks)

    # noinspection PyMethodMayBeStatic
    def create_pool(self):
        pool = ThreadPool(cpu_count())
        return pool

    def test_hexdigest(self):
        # Check basic 1-D array
        z = self.create_array(shape=(1050,), chunks=100, dtype='<i4')
        assert 'f710da18d45d38d4aaf2afd7fb822fdd73d02957' == z.hexdigest()

        # Check basic 1-D array with different type
        z = self.create_array(shape=(1050,), chunks=100, dtype='<f4')
        assert '1437428e69754b1e1a38bd7fc9e43669577620db' == z.hexdigest()

        # Check basic 2-D array
        z = self.create_array(shape=(20, 35,), chunks=10, dtype='<i4')
        assert '6c530b6b9d73e108cc5ee7b6be3d552cc994bdbe' == z.hexdigest()

        # Check basic 1-D array with some data
        z = self.create_array(shape=(1050,), chunks=100, dtype='<i4')
        z[200:400] = np.arange(200, 400, dtype='i4')
        assert '4c0a76fb1222498e09dcd92f7f9221d6cea8b40e' == z.hexdigest()

        # Check basic 1-D array with attributes
        z = self.create_array(shape=(1050,), chunks=100, dtype='<i4')
        z.attrs['foo'] = 'bar'
        assert '05b0663ffe1785f38d3a459dec17e57a18f254af' == z.hexdigest()


class TestArrayWithProcessSynchronizer(TestArray, MixinArraySyncTests):

    def create_array(self, read_only=False, **kwargs):
        path = tempfile.mkdtemp()
        atexit.register(atexit_rmtree, path)
        store = DirectoryStore(path)
        cache_metadata = kwargs.pop('cache_metadata', False)
        cache_attrs = kwargs.pop('cache_attrs', False)
        write_empty_chunks = kwargs.pop('write_empty_chunks', True)
        init_array(store, **kwargs)
        sync_path = tempfile.mkdtemp()
        atexit.register(atexit_rmtree, sync_path)
        synchronizer = ProcessSynchronizer(sync_path)
        return Array(store, synchronizer=synchronizer, read_only=read_only,
                     cache_metadata=cache_metadata, cache_attrs=cache_attrs,
                     write_empty_chunks=write_empty_chunks)

    # noinspection PyMethodMayBeStatic
    def create_pool(self):
        pool = ProcessPool(processes=cpu_count())
        return pool

    def test_hexdigest(self):
        # Check basic 1-D array
        z = self.create_array(shape=(1050,), chunks=100, dtype='<i4')
        assert 'f710da18d45d38d4aaf2afd7fb822fdd73d02957' == z.hexdigest()

        # Check basic 1-D array with different type
        z = self.create_array(shape=(1050,), chunks=100, dtype='<f4')
        assert '1437428e69754b1e1a38bd7fc9e43669577620db' == z.hexdigest()

        # Check basic 2-D array
        z = self.create_array(shape=(20, 35,), chunks=10, dtype='<i4')
        assert '6c530b6b9d73e108cc5ee7b6be3d552cc994bdbe' == z.hexdigest()

        # Check basic 1-D array with some data
        z = self.create_array(shape=(1050,), chunks=100, dtype='<i4')
        z[200:400] = np.arange(200, 400, dtype='i4')
        assert '4c0a76fb1222498e09dcd92f7f9221d6cea8b40e' == z.hexdigest()

        # Check basic 1-D array with attributes
        z = self.create_array(shape=(1050,), chunks=100, dtype='<i4')
        z.attrs['foo'] = 'bar'
        assert '05b0663ffe1785f38d3a459dec17e57a18f254af' == z.hexdigest()

    def test_object_arrays_danger(self):
        # skip this one, metadata get reloaded in each process
        pass


def _create_group(arg):
    g, name = arg
    h = g.create_group(name)
    return h.name


def _require_group(arg):
    g, name = arg
    h = g.require_group(name)
    return h.name


class MixinGroupSyncTests:

    def test_parallel_create_group(self):

        # setup
        g = self.create_group()
        pool = self.create_pool()

        # parallel create group
        n = 100
        results = list(pool.map(
            _create_group,
            zip([g] * n, [str(i) for i in range(n)]),
            chunksize=1
        ))
        assert n == len(results)
        pool.close()
        pool.terminate()

        assert n == len(g)

        pool.terminate()

    def test_parallel_require_group(self):

        # setup
        g = self.create_group()
        pool = self.create_pool()

        # parallel require group
        n = 100
        results = list(pool.map(
            _require_group,
            zip([g] * n, [str(i//10) for i in range(n)]),
            chunksize=1
        ))
        assert n == len(results)
        pool.close()
        pool.terminate()

        assert n//10 == len(g)

        pool.terminate()


class TestGroupWithThreadSynchronizer(TestGroup, MixinGroupSyncTests):

    def create_group(self, store=None, path=None, read_only=False,
                     chunk_store=None, synchronizer=None):
        if store is None:
            store, chunk_store = self.create_store()
        init_group(store, path=path, chunk_store=chunk_store)
        synchronizer = ThreadSynchronizer()
        g = Group(store, path=path, read_only=read_only,
                  chunk_store=chunk_store, synchronizer=synchronizer)
        return g

    # noinspection PyMethodMayBeStatic
    def create_pool(self):
        pool = ThreadPool(cpu_count())
        return pool

    def test_synchronizer_property(self):
        g = self.create_group()
        assert isinstance(g.synchronizer, ThreadSynchronizer)


class TestGroupWithProcessSynchronizer(TestGroup, MixinGroupSyncTests):

    def create_store(self):
        path = tempfile.mkdtemp()
        atexit.register(atexit_rmtree, path)
        store = DirectoryStore(path)
        return store, None

    def create_group(self, store=None, path=None, read_only=False,
                     chunk_store=None, synchronizer=None):
        if store is None:
            store, chunk_store = self.create_store()
        init_group(store, path=path, chunk_store=chunk_store)
        sync_path = tempfile.mkdtemp()
        atexit.register(atexit_rmtree, sync_path)
        synchronizer = ProcessSynchronizer(sync_path)
        g = Group(store, path=path, read_only=read_only,
                  synchronizer=synchronizer, chunk_store=chunk_store)
        return g

    # noinspection PyMethodMayBeStatic
    def create_pool(self):
        pool = ProcessPool(processes=cpu_count())
        return pool

    def test_synchronizer_property(self):
        g = self.create_group()
        assert isinstance(g.synchronizer, ProcessSynchronizer)