"""
Test suite.
"""

import math
import unittest
import torch

from continual.src.utils import models, 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"]


def magnitude(value):
    """
    Return the magnitude of a positive number.
    """
    if value < 0:
        raise ValueError
    if value == 0:
        return 0
    else:
        return int(math.floor(math.log10(value)))


class TestModelMethods(unittest.TestCase):
    """
    Model definition tests.
    """

    def test_modeloutputshape(self):
        """
        Testing model produces correct shape of output for variety of input sizes.
        """
        for batch_size in BATCH_SIZES:
            for seq_len in SEQ_LENS:
                for n_vars in N_VARS:
                    for n_classes in N_CLASSES:
                        for n_layers in N_LAYERS:
                            for hidden_size in HIDDEN_SIZES:
                                batch = torch.randn(batch_size, seq_len, n_vars)
                                simple_models = models.MODELS.values()
                                for model in simple_models:
                                    model = model(
                                        seq_len=seq_len,
                                        n_channels=n_vars,
                                        hidden_dim=hidden_size,
                                        output_size=n_classes,
                                        n_layers=n_layers,
                                    )
                                    # Set in eval mode to avoid batch-norm error when subtracting mean from val training on 1 datapoint
                                    model.eval()
                                    output = model(batch)
                                    expected_shape = torch.Size([batch_size, n_classes])
                                    self.assertEqual(output.shape, expected_shape)

    def ttest_modelcapacity(self):
        """
        JA: Need to update given parameterisation of model structure.

        Testing different models have same order of magnitude of parameters.
        """
        for seq_len in SEQ_LENS:
            for n_vars in N_VARS:
                for n_classes in N_CLASSES:
                    simple_models = models.MODELS.values()
                    n_params = [
                        sum(
                            p.numel()
                            for p in m(
                                seq_len=seq_len,
                                n_channels=n_vars,
                                output_size=n_classes,
                            ).parameters()
                            if p.requires_grad
                        )
                        for m in simple_models
                    ]
                    param_magnitudes = [magnitude(p) for p in n_params]
                    # RNN/LSTM order bigger
                    self.assertTrue(max(param_magnitudes) - min(param_magnitudes) <= 1)

    # JA: Implement test to check params passed by config actually change model structure.


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]))


if __name__ == "__main__":
    unittest.main()