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

from core.ensemble.model_lib import build_models
import numpy as np


class Model_lib_test(parameterized.TestCase):
    @parameterized.named_parameters(
        [(ens_type, ens_type) for ens_type in \
         ['vanilla', 'deep_ensemble', 'weighted_ensemble', 'mc_dropout', 'bayesian']])
    def test_dnn_graph(self, ensemble_type):
        x = np.array([[1., 2., 3.], [4., 5., 6.]])
        if ensemble_type is not 'mc_dropout':
            build_models(x, 'dnn', ensemble_type, input_dim=3)
        else:
            build_models(x, 'dnn', ensemble_type, input_dim=3, use_mc_dropout=True)

    @parameterized.named_parameters(
        [(ens_type, ens_type) for ens_type in \
         ['vanilla', 'deep_ensemble', 'weighted_ensemble', 'mc_dropout', 'bayesian']])
    def test_text_cnn_graph(self, ensemble_type):
        x = np.random.randint(0, 256, (2, 16))
        if ensemble_type is not 'mc_dropout':
            out = build_models(x, 'text_cnn', ensemble_type, input_dim=5)
        else:
            out = build_models(x, 'text_cnn', ensemble_type, input_dim=5, use_mc_dropout=True)

        self.assertTrue(out.shape == (2, 1))

    @parameterized.named_parameters(
        [(ens_type, ens_type) for ens_type in \
         ['vanilla', 'deep_ensemble', 'weighted_ensemble', 'mc_dropout', 'bayesian']])
    def test_multimodalitynn(self, ensemble_type):
        _x = np.array([[1., 2., 3.], [4., 5., 6.]])
        x = [_x, _x, _x, _x, _x]
        if ensemble_type is not 'mc_dropout':
            out = build_models(x, 'multimodalitynn', ensemble_type, input_dim=[3, 3, 3, 3, 3])
        else:
            out = build_models(x, 'multimodalitynn', ensemble_type, input_dim=[3, 3, 3, 3, 3], use_mc_dropout=True)
        self.assertTrue(out.shape == (2, 1))

    @parameterized.named_parameters(
        [(ens_type, ens_type) for ens_type in \
    ['vanilla', 'deep_ensemble', 'weighted_ensemble', 'mc_dropout', 'bayesian']])
    def test_droidectc(self, ensemble_type):
        _x = np.random.randint(0, 10000, [2, 1000])
        if ensemble_type is not 'mc_dropout':
            out = build_models(_x, 'droidectc', ensemble_type)
        else:
            out = build_models(_x, 'droidectc', ensemble_type, use_mc_dropout=True)
        self.assertTrue(out.shape == (2, 1))


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