import unittest

import numpy as np

from chainer import testing
from chainercv.transforms import random_expand


class TestRandomExpand(unittest.TestCase):

    def test_random_expand(self):
        img = np.random.uniform(-1, 1, size=(3, 64, 32))

        out = random_expand(img)

        out = random_expand(img, max_ratio=1)
        np.testing.assert_equal(out, img)

        out, ratio, x_offset, y_offset = random_expand(
            img, max_ratio=4, return_params=True)
        np.testing.assert_equal(
            out[:, y_offset:y_offset + 64, x_offset:x_offset + 32], img)
        self.assertGreaterEqual(ratio, 1)
        self.assertLessEqual(ratio, 4)
        self.assertEqual(out.shape[1], int(64 * ratio))
        self.assertEqual(out.shape[2], int(32 * ratio))

        out = random_expand(img, max_ratio=2)

    def test_random_expand_fill(self):
        img = np.random.uniform(-1, 1, size=(3, 64, 32))

        fills = (128, (104, 117, 123), np.random.uniform(255, size=3))

        for fill in fills:
            while True:
                out, _, x_offset, y_offset = random_expand(
                    img, fill=fill, return_params=True)
                if x_offset > 0 or y_offset > 0:
                    break

            if isinstance(fill, int):
                np.testing.assert_equal(
                    out[:, 0, 0], (fill, fill, fill))
            else:
                np.testing.assert_equal(
                    out[:, 0, 0], fill)


testing.run_module(__name__, __file__)

import unittest

import numpy as np

from chainer import testing
from chainercv.transforms import random_expand


class TestRandomExpand(unittest.TestCase):

    def test_random_expand(self):
        img = np.random.uniform(-1, 1, size=(3, 64, 32))

        out = random_expand(img)

        out = random_expand(img, max_ratio=1)
        np.testing.assert_equal(out, img)

        out, ratio, x_offset, y_offset = random_expand(
            img, max_ratio=4, return_params=True)
        np.testing.assert_equal(
            out[:, y_offset:y_offset + 64, x_offset:x_offset + 32], img)
        self.assertGreaterEqual(ratio, 1)
        self.assertLessEqual(ratio, 4)
        self.assertEqual(out.shape[1], int(64 * ratio))
        self.assertEqual(out.shape[2], int(32 * ratio))

        out = random_expand(img, max_ratio=2)


@testing.parameterize(
    {'fill': 128},
    {'fill': (104, 117, 123)},
    {'fill':  np.random.uniform(255, size=3)},
)
class TestRandomExpandFill(unittest.TestCase):

    def test_random_expand_fill(self):
        img = np.random.uniform(-1, 1, size=(3, 64, 32))

        while True:
            out, _, x_offset, y_offset = random_expand(
                img, fill=self.fill, return_params=True)
            if x_offset > 0 or y_offset > 0:
                break

        if isinstance(self.fill, int):
            np.testing.assert_equal(
                out[:, 0, 0], (self.fill,) * 3)
        else:
            np.testing.assert_equal(
                out[:, 0, 0], self.fill)


testing.run_module(__name__, __file__)

use testing.parameterize