import unittest
import dace
import numpy as np
from dace.transformation.dataflow import MapTiling, OutLocalStorage

N = dace.symbol('N')


@dace.program
def arange():
    out = np.ndarray([N], np.int32)
    for i in dace.map[0:N]:
        with dace.tasklet:
            o >> out[i]
            o = i
    return out


class LocalStorageTests(unittest.TestCase):
    def test_even(self):
        sdfg = arange.to_sdfg()
        sdfg.apply_transformations([MapTiling, OutLocalStorage],
                                   options=[{
                                       'tile_sizes': [8]
                                   }, {}])
        self.assertTrue(
            np.array_equal(sdfg(N=16), np.arange(16, dtype=np.int32)))

    def test_uneven(self):
        # For testing uneven decomposition, use longer buffer and ensure
        # it's not filled over
        output = np.ones(20, np.int32)
        sdfg = arange.to_sdfg()
        sdfg.apply_transformations([MapTiling, OutLocalStorage],
                                   options=[{
                                       'tile_sizes': [5]
                                   }, {}])
        dace.propagate_memlets_sdfg(sdfg)
        sdfg(N=16, __return=output)
        self.assertTrue(
            np.array_equal(output[:16], np.arange(16, dtype=np.int32)))
        self.assertTrue(np.array_equal(output[16:], np.ones(4, np.int32)))


if __name__ == '__main__':
    unittest.main()

[((6, 4, 6, 20), 'dace.symbol', 'dace.symbol', ({(6, 16, 6, 19): '"""N"""'}, {}), "('N')", False, 'import dace\n'), ((11, 10, 11, 35), 'numpy.ndarray', 'np.ndarray', ({(11, 21, 11, 24): '[N]', (11, 26, 11, 34): 'np.int32'}, {}), '([N], np.int32)', True, 'import numpy as np\n'), ((46, 4, 46, 19), 'unittest.main', 'unittest.main', ({}, {}), '()', False, 'import unittest\n'), ((32, 17, 32, 38), 'numpy.ones', 'np.ones', ({(32, 25, 32, 27): '20', (32, 29, 32, 37): 'np.int32'}, {}), '(20, np.int32)', True, 'import numpy as np\n'), ((38, 8, 38, 41), 'dace.propagate_memlets_sdfg', 'dace.propagate_memlets_sdfg', ({(38, 36, 38, 40): 'sdfg'}, {}), '(sdfg)', False, 'import dace\n'), ((27, 39, 27, 68), 'numpy.arange', 'np.arange', (), '', True, 'import numpy as np\n'), ((41, 40, 41, 69), 'numpy.arange', 'np.arange', (), '', True, 'import numpy as np\n'), ((42, 52, 42, 72), 'numpy.ones', 'np.ones', ({(42, 60, 42, 61): '(4)', (42, 63, 42, 71): 'np.int32'}, {}), '(4, np.int32)', True, 'import numpy as np\n')]