import unittest

import numpy as np

from sigpy import block, config

if config.cupy_enabled:
    import cupy as cp

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


class TestInterp(unittest.TestCase):
    def test_array_to_blocks(self):
        xps = [np]
        if config.cupy_enabled:
            xps.append(cp)

        for xp in xps:
            for dtype in [np.float32, np.complex64]:
                for ndim in [1, 2, 3]:
                    with self.subTest(xp=xp, dtype=dtype, ndim=ndim):
                        input = xp.array(
                            [0, 1, 2, 3, 4, 5], dtype=dtype
                        ).reshape([6] + [1] * (ndim - 1))

                        blk_shape = [1] + [1] * (ndim - 1)
                        blk_strides = [1] + [1] * (ndim - 1)
                        output = xp.array(
                            [[0], [1], [2], [3], [4], [5]], dtype=dtype
                        ).reshape(
                            [6] + [1] * (ndim - 1) + [1] + [1] * (ndim - 1)
                        )
                        xp.testing.assert_allclose(
                            output,
                            block.array_to_blocks(
                                input, blk_shape, blk_strides
                            ),
                        )

                        blk_shape = [2] + [1] * (ndim - 1)
                        blk_strides = [1] + [1] * (ndim - 1)
                        output = xp.array(
                            [[0, 1], [1, 2], [2, 3], [3, 4], [4, 5]],
                            dtype=dtype,
                        ).reshape(
                            [5] + [1] * (ndim - 1) + [2] + [1] * (ndim - 1)
                        )
                        xp.testing.assert_allclose(
                            output,
                            block.array_to_blocks(
                                input, blk_shape, blk_strides
                            ),
                        )

                        blk_shape = [2] + [1] * (ndim - 1)
                        blk_strides = [2] + [1] * (ndim - 1)
                        output = xp.array(
                            [[0, 1], [2, 3], [4, 5]], dtype=dtype
                        ).reshape(
                            [3] + [1] * (ndim - 1) + [2] + [1] * (ndim - 1)
                        )
                        xp.testing.assert_allclose(
                            output,
                            block.array_to_blocks(
                                input, blk_shape, blk_strides
                            ),
                        )

                        blk_shape = [3] + [1] * (ndim - 1)
                        blk_strides = [2] + [1] * (ndim - 1)
                        output = xp.array(
                            [[0, 1, 2], [2, 3, 4]], dtype=dtype
                        ).reshape(
                            [2] + [1] * (ndim - 1) + [3] + [1] * (ndim - 1)
                        )
                        xp.testing.assert_allclose(
                            output,
                            block.array_to_blocks(
                                input, blk_shape, blk_strides
                            ),
                        )

    def test_blocks_to_array(self):
        xps = [np]
        if config.cupy_enabled:
            xps.append(cp)

        for xp in xps:
            for dtype in [np.float32, np.complex64]:
                for ndim in [1, 2, 3]:
                    with self.subTest(xp=xp, dtype=dtype, ndim=ndim):
                        shape = [6] + [1] * (ndim - 1)

                        blk_shape = [1] + [1] * (ndim - 1)
                        blk_strides = [1] + [1] * (ndim - 1)
                        input = xp.array(
                            [[0], [1], [2], [3], [4], [5]], dtype=dtype
                        ).reshape(
                            [6] + [1] * (ndim - 1) + [1] + [1] * (ndim - 1)
                        )
                        output = xp.array(
                            [0, 1, 2, 3, 4, 5], dtype=dtype
                        ).reshape([6] + [1] * (ndim - 1))
                        xp.testing.assert_allclose(
                            output,
                            block.blocks_to_array(
                                input, shape, blk_shape, blk_strides
                            ),
                        )

                        blk_shape = [2] + [1] * (ndim - 1)
                        blk_strides = [1] + [1] * (ndim - 1)
                        input = xp.array(
                            [[0, 1], [1, 2], [2, 3], [3, 4], [4, 5]],
                            dtype=dtype,
                        ).reshape(
                            [5] + [1] * (ndim - 1) + [2] + [1] * (ndim - 1)
                        )
                        output = xp.array(
                            [0, 2, 4, 6, 8, 5], dtype=dtype
                        ).reshape([6] + [1] * (ndim - 1))
                        xp.testing.assert_allclose(
                            output,
                            block.blocks_to_array(
                                input, shape, blk_shape, blk_strides
                            ),
                        )

                        blk_shape = [2] + [1] * (ndim - 1)
                        blk_strides = [2] + [1] * (ndim - 1)
                        input = xp.array(
                            [[0, 1], [2, 3], [4, 5]], dtype=dtype
                        ).reshape(
                            [3] + [1] * (ndim - 1) + [2] + [1] * (ndim - 1)
                        )
                        output = xp.array(
                            [0, 1, 2, 3, 4, 5], dtype=dtype
                        ).reshape([6] + [1] * (ndim - 1))
                        xp.testing.assert_allclose(
                            output,
                            block.blocks_to_array(
                                input, shape, blk_shape, blk_strides
                            ),
                        )

                        blk_shape = [3] + [1] * (ndim - 1)
                        blk_strides = [2] + [1] * (ndim - 1)
                        input = xp.array(
                            [[0, 1, 2], [2, 3, 4]], dtype=dtype
                        ).reshape(
                            [2] + [1] * (ndim - 1) + [3] + [1] * (ndim - 1)
                        )
                        output = xp.array(
                            [0, 1, 4, 3, 4, 0], dtype=dtype
                        ).reshape([6] + [1] * (ndim - 1))
                        xp.testing.assert_allclose(
                            output,
                            block.blocks_to_array(
                                input, shape, blk_shape, blk_strides
                            ),
                        )
