"""

Test suite.

"""

import unittest

import itertools

import torch

from continual.src.utils import data_processing

BATCH_SIZES = (1, 10, 100)

SEQ_LENS = (4, 12, 48)

N_VARS = (2, 10, 100)

N_CLASSES = (2, 10)

N_LAYERS = (1, 2, 3, 4)

HIDDEN_SIZES = (32, 64, 128)

DEMOGRAPHICS = [

    "age",

    "gender",

    "ethnicity",

    "region",

    "time_year",

    "time_season",

    "time_month",

]

OUTCOMES = ["ARF", "shock", "mortality"]

DATASETS = ["MIMIC", "eICU"]

class TestDataLoadingMethods(unittest.TestCase):

    """

    Data loading tests.

    """

    def test_modalfeatvalfromseq(self):

        """

        Test that mode of correct dim is returned.

        """

        for n_samples in BATCH_SIZES:

            for seq_len in SEQ_LENS:

                for n_feats in N_VARS:

                    for i in range(n_feats):

                        sim_data = (

                            torch.randint(0, 1, (n_samples, seq_len, n_feats))

                            .clone()

                            .detach()

                            .numpy()

                        )

                        modes = data_processing.get_modes(sim_data, feat=i)

                        self.assertEqual(modes.shape, torch.Size([n_samples]))

# CL task split tests

class TestCLConstructionMethods(unittest.TestCase):

    """

    Test construction of Continual Learning task splits.

    """

    def ttest_taskidsnonoverlap(self):

        for dataset in DATASETS:

            for experiment in OUTCOMES:

                for demographic in DEMOGRAPHICS:

                    # JA: implement

                    tasks = data_processing.load_data(dataset, demographic, experiment)

                    for pair in itertools.combinations(tasks, repeat=2):

                        self.assertTrue(pair[0][:, 0].intersection(pair[0][:, 0]) == {})

    def ttest_tasktargets(self):

        for dataset in DATASETS:

            for experiment in OUTCOMES:

                for demographic in DEMOGRAPHICS:

                    # JA: implement

                    tasks = data_processing.load_data(dataset, demographic, experiment)

                    for task in tasks:

                        self.assertTrue(len(task[:, -1].unique()) == 2)

if __name__ == "__main__":

    unittest.main()