# encoding: utf-8
# ---------------------------------------------------------------------------
#  Copyright (C) 2008-2014, IPython Development Team and Enthought, Inc.
#  Distributed under the terms of the BSD License.  See COPYING.rst.
# ---------------------------------------------------------------------------

import unittest
from random import randrange

from distarray.externals.six.moves import range

from distarray.dist.context import Context
from distarray.dist import maps as client_map


class TestClientMap(unittest.TestCase):

    def setUp(self):
        self.ctx = Context()

    def tearDown(self):
        self.ctx.close()

    def test_2D_bn(self):
        nrows, ncols = 31, 53
        cm = client_map.Distribution.from_shape(self.ctx, (nrows, ncols),
                                                {0: 'b'}, (4, 1))
        chunksize = (nrows // 4) + 1
        for _ in range(100):
            r, c = randrange(nrows), randrange(ncols)
            rank = r // chunksize
            self.assertSequenceEqual(cm.owning_ranks((r,c)), [rank])

    def test_2D_bb(self):
        nrows, ncols = 3, 5
        nprocs_per_dim = 2
        cm = client_map.Distribution.from_shape(
                self.ctx, (nrows, ncols), ('b', 'b'),
                (nprocs_per_dim, nprocs_per_dim))
        row_chunks = nrows // nprocs_per_dim + 1
        col_chunks = ncols // nprocs_per_dim + 1
        for r in range(nrows):
            for c in range(ncols):
                rank = (r // row_chunks) * nprocs_per_dim + (c // col_chunks)
                actual = cm.owning_ranks((r,c))
                self.assertSequenceEqual(actual, [rank])

    def test_2D_cc(self):
        nrows, ncols = 3, 5
        nprocs_per_dim = 2
        cm = client_map.Distribution.from_shape(
                self.ctx, (nrows, ncols), ('c', 'c'),
                (nprocs_per_dim, nprocs_per_dim))
        for r in range(nrows):
            for c in range(ncols):
                rank = (r % nprocs_per_dim) * nprocs_per_dim + (c % nprocs_per_dim)
                actual = cm.owning_ranks((r,c))
                self.assertSequenceEqual(actual, [rank])


    def test_is_compatible(self):
        nr, nc, nd = 10**5, 10**6, 10**4

        cm0 = client_map.Distribution.from_shape(
                 self.ctx, (nr, nc, nd), ('b', 'c', 'n'))
        self.assertTrue(cm0.is_compatible(cm0))

        cm1 = client_map.Distribution.from_shape(
                 self.ctx, (nr, nc, nd), ('b', 'c', 'n'))
        self.assertTrue(cm1.is_compatible(cm1))

        self.assertTrue(cm0.is_compatible(cm1))
        self.assertTrue(cm1.is_compatible(cm0))
        
        nr -= 1; nc -= 1; nd -= 1

        cm2 = client_map.Distribution.from_shape(
                 self.ctx, (nr, nc, nd), ('b', 'c', 'n'))

        self.assertFalse(cm1.is_compatible(cm2))
        self.assertFalse(cm2.is_compatible(cm1))

    def test_reduce(self):
        nr, nc, nd = 10**5, 10**6, 10**4

        dist = client_map.Distribution.from_shape(
                 self.ctx, (nr, nc, nd), ('b', 'c', 'n'), grid_shape=(2, 2, 1))

        new_dist0 = dist.reduce(axis=0)
        self.assertEquals(new_dist0.dist, ('c', 'n'))
        self.assertEquals(new_dist0.shape, (nc, nd))
        self.assertEquals(new_dist0.grid_shape, dist.grid_shape[1:])
        self.assertLess(set(new_dist0.targets), set(dist.targets))

        new_dist1 = dist.reduce(axis=1)
        self.assertEquals(new_dist1.dist, ('b', 'n'))
        self.assertEquals(new_dist1.shape, (nr, nd))
        self.assertEquals(new_dist1.grid_shape, dist.grid_shape[:1]+dist.grid_shape[2:])
        self.assertLess(set(new_dist1.targets), set(dist.targets))

        new_dist2 = dist.reduce(axis=2)
        self.assertEquals(new_dist2.dist, ('b', 'c'))
        self.assertEquals(new_dist2.shape, (nr, nc))
        self.assertEquals(new_dist2.grid_shape, dist.grid_shape[:-1])
        self.assertEquals(set(new_dist2.targets), set(dist.targets))

@@ -3107,17 +3107,19 @@
 ce(axis=
-0
+%5B0%5D
 )%0A      
@@ -3393,17 +3393,19 @@
 ce(axis=
-1
+%5B1%5D
 )%0A      
@@ -3699,17 +3699,19 @@
 ce(axis=
-2
+%5B2%5D
 )%0A      

Fix axis argument to distribution.reduce()