from absl.testing import absltest
from absl.testing import parameterized
import tempfile

import numpy as np
import tensorflow as tf
from sklearn.datasets import load_breast_cancer

from core.ensemble.vanilla import Vanilla
from core.ensemble.dataset_lib import build_dataset_from_numerical_data


class MyTestCaseVanilla(parameterized.TestCase):
    def setUp(self):
        self.x_np, self.y_np = load_breast_cancer(return_X_y=True)
        self.train_dataset_v1 = build_dataset_from_numerical_data((self.x_np, self.y_np))
        self.val_dataset_v1 = build_dataset_from_numerical_data((self.x_np, self.y_np))

    def test_vanilla_dnn(self):
        with tempfile.TemporaryDirectory() as output_dir:
            vanilla = Vanilla(architecture_type='dnn',
                              model_directory=output_dir)
            vanilla.fit(self.train_dataset_v1, self.val_dataset_v1, input_dim=self.x_np.shape[1])

            res = vanilla.predict(self.x_np)
            self.assertEqual(vanilla.get_n_members(), vanilla.n_members)
            self.assertTrue(res.shape == (self.x_np.shape[0], vanilla.n_members, 1))

            vanilla.evaluate(self.x_np, self.y_np)

    def test_vanilla_textcnn(self):
        with tempfile.TemporaryDirectory() as output_dir:
            vanilla = Vanilla(architecture_type='text_cnn',
                              model_directory=output_dir)
            x = np.random.randint(0, 256, (10, 10))
            y = np.random.choice(2, 10)
            train_dataset = build_dataset_from_numerical_data((x,y))
            val_dataset = build_dataset_from_numerical_data((x,y))
            vanilla.fit(train_dataset, val_dataset)
            res = vanilla.predict(x)
            self.assertTrue(res.shape == (x.shape[0], vanilla.n_members, 1))
            vanilla.evaluate(x, y)

    def test_vanilla_multimodalitynn(self):
        with tempfile.TemporaryDirectory() as output_dir:
            vanilla = Vanilla(architecture_type='multimodalitynn',
                              model_directory=output_dir)
            x = [self.x_np] * 5
            train_data = build_dataset_from_numerical_data(tuple(x))
            train_y = build_dataset_from_numerical_data(self.y_np)
            train_dataset = tf.data.Dataset.zip((train_data, train_y))
            val_data = build_dataset_from_numerical_data(tuple(x))
            val_y = build_dataset_from_numerical_data(self.y_np)
            val_dataset = tf.data.Dataset.zip((val_data, val_y))
            vanilla.fit(train_dataset, val_dataset, input_dim=[self.x_np.shape[1]]*5)
            res = vanilla.predict(x)
            self.assertTrue(res.shape == (self.x_np.shape[0], vanilla.n_members, 1))
            vanilla.evaluate(x, self.y_np)

    def test_vanilla_r2d2(self):
        with tempfile.TemporaryDirectory() as output_dir:
            vanilla = Vanilla(architecture_type='r2d2',
                              model_directory=output_dir)
            x = np.random.uniform(0., 1., size=(10, 299, 299, 3))
            y = np.random.choice(2, 10)
            train_dataset = build_dataset_from_numerical_data((x, y))
            val_dataset = build_dataset_from_numerical_data((x, y))
            vanilla.fit(train_dataset, val_dataset, input_dim=(299, 299, 3))
            res = vanilla.predict(x)
            self.assertTrue(res.shape == (x.shape[0], vanilla.n_members, 1))
            vanilla.evaluate(x, y)

    def test_vanilla_droidectc(self):
        with tempfile.TemporaryDirectory() as output_dir:
            vanilla = Vanilla(architecture_type='droidectc',
                              model_directory=output_dir)

            x = np.random.randint(0, 10000, size=(10, 1000))
            y = np.random.choice(2, 10)
            train_dataset = build_dataset_from_numerical_data((x, y))
            val_dataset = build_dataset_from_numerical_data((x, y))
            vanilla.fit(train_dataset, val_dataset)
            res = vanilla.predict(x)
            self.assertTrue(res.shape == (x.shape[0], vanilla.n_members, 1))
            vanilla.evaluate(x, y)


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