# Copyright (c) Facebook, Inc. and its affiliates.
import unittest
from collections import Counter

import numpy as np
import torch
from mmf.common.registry import registry
from mmf.datasets.multi_dataset_loader import MultiDatasetLoader
from omegaconf import OmegaConf
from torch.utils.data import DataLoader

from ..test_utils import NumbersDataset


class TestMultiDatasetLoader(unittest.TestCase):
    def setUp(self):
        torch.manual_seed(1234)
        np.random.seed(1234)
        self.multi_dataset = MultiDatasetLoader()
        self.multi_dataset._num_datasets = 3
        self.multi_dataset.current_index = 0
        numbers_dataset_a = NumbersDataset(4, "a")
        numbers_dataset_b = NumbersDataset(40, "b")
        numbers_dataset_c = NumbersDataset(4000, "c")
        self.multi_dataset.dataset_list = ["a", "b", "c"]
        self.multi_dataset._loaders = {
            "a": self._get_dataloader(numbers_dataset_a),
            "b": self._get_dataloader(numbers_dataset_b),
            "c": self._get_dataloader(numbers_dataset_c),
        }
        self.original_config = registry.get("config")
        registry.register(
            "config",
            OmegaConf.create(
                {
                    "training": {
                        "dataset_size_proportional_sampling": True,
                        "max_epochs": None,
                    },
                    "multitasking": {"enabled": False},
                }
            ),
        )
        self.multi_dataset._per_dataset_lengths = [4, 40, 4000]
        self.multi_dataset._total_length = 4044
        self.multi_dataset._total_length = sum(self.multi_dataset._per_dataset_lengths)

    def tearDown(self):
        registry.register("config", self.original_config)

    def _get_dataloader(self, dataset):
        return DataLoader(dataset=dataset, batch_size=4, num_workers=0)

    def test_proportional_sampling(self):
        self.multi_dataset._infer_dataset_probabilities()

        count = 0
        count_c = 0
        for batch in self.multi_dataset:
            batch = self.multi_dataset.prepare_batch(batch)
            if "c" in batch:
                count_c += 1
            count += 1
            if count == 100:
                break

        # Expect more than 95 c's at least as the len for c is very high
        self.assertTrue(count_c >= 98)

        count = 0
        count_epoch = 0
        counter = Counter()
        for _ in range(1):
            for batch in self.multi_dataset:
                batch = self.multi_dataset.prepare_batch(batch)
                counter[list(batch.keys())[0]] += 1
                count += 1
            count_epoch += 1
        # Expect epoch to be completed
        self.assertEqual(count_epoch, 1)
        # Expect each dataset to be full iterated
        self.assertEqual(count, self.multi_dataset._total_length // 4)
        self.assertEqual(counter, Counter({"a": 1, "b": 10, "c": 1000}))

    def test_equal_sampling(self):
        registry.get("config").training.dataset_size_proportional_sampling = False
        self.multi_dataset._infer_dataset_probabilities()
        count = 0
        count_c = 0
        for batch in self.multi_dataset:
            batch = self.multi_dataset.prepare_batch(batch)
            if "c" in batch:
                count_c += 1
            count += 1
            if count == 100:
                break

        self.assertTrue(count_c <= 34)

        # Epoch will never finish for this case, so test upto proportional sampling's
        # epoch length + some extra
        for batch in self.multi_dataset:
            batch = self.multi_dataset.prepare_batch(batch)
            count += 1
            if count > self.multi_dataset._total_length // 4 + 100:
                break

        # The test should reach at this stage and should not be finished at
        # epoch length
        self.assertTrue(count > self.multi_dataset._total_length // 4 + 100)
