""" Compose multiple datasets in a single loader. """

import numpy as np
from copy import deepcopy
from torch.utils.data import Dataset

from .wireframe_dataset import WireframeDataset
from .holicity_dataset import HolicityDataset


class MergeDataset(Dataset):
    def __init__(self, mode, config=None):
        super(MergeDataset, self).__init__()
        # Initialize the datasets
        self._datasets = []
        spec_config = deepcopy(config)
        for i, d in enumerate(config["datasets"]):
            spec_config["dataset_name"] = d
            spec_config["gt_source_train"] = config["gt_source_train"][i]
            spec_config["gt_source_test"] = config["gt_source_test"][i]
            if d == "wireframe":
                self._datasets.append(WireframeDataset(mode, spec_config))
            elif d == "holicity":
                spec_config["train_split"] = config["train_splits"][i]
                self._datasets.append(HolicityDataset(mode, spec_config))
            else:
                raise ValueError("Unknown dataset: " + d)

        self._weights = config["weights"]

    def __getitem__(self, item):
        dataset = self._datasets[
            np.random.choice(range(len(self._datasets)), p=self._weights)
        ]
        return dataset[np.random.randint(len(dataset))]

    def __len__(self):
        return np.sum([len(d) for d in self._datasets])