from __future__ import division, print_function

import unittest

import numpy as np
from rvseg import dataset

class TestDataset(unittest.TestCase):
    def test_generator(self):
        self._test_generator(mask='inner')
        self._test_generator(mask='outer')
        self._test_generator(mask='both')

    def test_no_validation(self):
        self._test_no_validation(mask='inner')
        self._test_no_validation(mask='outer')
        self._test_no_validation(mask='both')

    def _test_generator(self, mask):
        data_dir = "../test-assets/"
        batch_size = 2
        validation_split = 0.5
        # With a total of 3 training images, this split will create 1
        # training image and 2 validation images

        (train_generator, train_steps_per_epoch,
         val_generator, val_steps_per_epoch) = dataset.create_generators(
             data_dir, batch_size,
             validation_split=validation_split,
             mask=mask)

        self.assertEqual(train_steps_per_epoch, 1)
        self.assertEqual(val_steps_per_epoch, 1)

        classes = 3 if mask == 'both' else 2

        images, masks = next(train_generator)
        self.assertEqual(images.shape, (1, 216, 256, 1))
        self.assertEqual(masks.shape, (1, 216, 256, classes))

        images, masks = next(val_generator)
        self.assertEqual(images.shape, (2, 216, 256, 1))
        self.assertEqual(masks.shape, (2, 216, 256, classes))

    def _test_no_validation(self, mask):
        data_dir = "../test-assets/"
        batch_size = 2
        validation_split = 0.0

        (train_generator, train_steps_per_epoch,
         val_generator, val_steps_per_epoch) = dataset.create_generators(
             data_dir, batch_size,
             validation_split=validation_split,
             mask=mask)

        self.assertEqual(train_steps_per_epoch, 2)
        self.assertEqual(val_steps_per_epoch, 0)

        classes = 3 if mask == 'both' else 2

        # first 2 train images
        images, masks = next(train_generator)
        self.assertEqual(images.shape, (2, 216, 256, 1))
        self.assertEqual(masks.shape, (2, 216, 256, classes))

        # last train image (for total of 3)
        images, masks = next(train_generator)
        self.assertEqual(images.shape, (1, 216, 256, 1))
        self.assertEqual(masks.shape, (1, 216, 256, classes))

        # first 2 train images again
        images, masks = next(train_generator)
        self.assertEqual(images.shape, (2, 216, 256, 1))
        self.assertEqual(masks.shape, (2, 216, 256, classes))

        # validation generator should be nothing
        self.assertEqual(val_generator, None)


    def test_shuffle_train_val(self):
        # test shuffling of entire dataset prior to train-val split
        # (does not test shuffling within each epoch)
        data_dir = "../test-assets/"
        batch_size = 2
        validation_split = 0.5
        mask = "inner"
        classes = 2
        seed = 5               # random number seed

        # there should be 2 images in the validation set, and we'll check if
        # they always appear in the same order with a fixed seed
        image_list = []
        mask_list = []
        for i in range(10):
            _, _, val_generator, _ = dataset.create_generators(
                data_dir, batch_size, validation_split=validation_split,
                mask=mask, shuffle_train_val=True, shuffle=False, seed=seed,
                normalize_images=True)

            images, masks = next(val_generator)
            self.assertEqual(images.shape, (2, 216, 256, 1))
            self.assertEqual(masks.shape, (2, 216, 256, classes))

            # also check image normalization
            for image in images:
                self.assertAlmostEqual(np.mean(image), 0)
                self.assertAlmostEqual(np.std(image), 1, places=5)

            image_list.append(images[0])
            mask_list.append(masks[0])

        # first image/mask in each case should be the same
        image0 = image_list[0]
        for image in image_list[1:]:
            np.testing.assert_array_equal(image0, image)
        mask0 = mask_list[0]
        for mask in mask_list[1:]:
            np.testing.assert_array_equal(mask0, mask)

        # now test that things get shuffled if we don't specify a seed
        mask = "both"
        _, _, val_generator, _ = dataset.create_generators(
            data_dir, batch_size, validation_split=validation_split,
            mask=mask, shuffle_train_val=True, shuffle=False, seed=None,
            normalize_images=True)

        images, masks = next(val_generator)
        image0 = images[0]
        while 1:
            _, _, val_generator, _ = dataset.create_generators(
                data_dir, batch_size, validation_split=validation_split,
                mask=mask, shuffle_train_val=True, shuffle=True, seed=None,
                normalize_images=True)
            images, masks = next(val_generator)            
            try:
                np.testing.assert_array_equal(image0, images[0])
            except AssertionError:
                break           # break if arrays are differet (= success!)