from functools import partial

from kappadata.datasets.kd_dataset import KDDataset


class ModeWrapper(KDDataset):
    def __init__(self, dataset: KDDataset, mode: str, return_ctx: bool = False):
        super().__init__()
        self.dataset = dataset
        self.mode = mode
        self.return_ctx = return_ctx
        self.propagate_ctx = return_ctx

        self._getitem_fns = []
        self.items = mode.split(" ")
        for item in self.items:
            if item == "index":
                self._getitem_fns.append(self._getitem_index)
            elif item.startswith("ctx."):
                self.propagate_ctx = True
                ctx_key = item[len("ctx."):]
                self._getitem_fns.append(partial(self._getitem_from_ctx, ctx_key=ctx_key))
            else:
                fn_name = f"getitem_{item}"
                assert hasattr(self.dataset, fn_name), f"{type(self.dataset.root_dataset)} has no method getitem_{item}"
                self._getitem_fns.append(getattr(self.dataset, fn_name))

    @staticmethod
    def has_item(mode, item):
        return item in mode.split(" ")

    @staticmethod
    def get_item_index(mode, item):
        return mode.split(" ").index(item)

    @staticmethod
    def get_item(mode, item, batch):
        idx = ModeWrapper.get_item_index(mode=mode, item=item)
        return batch[idx]

    @staticmethod
    def set_item(mode, item, batch, value):
        idx = mode.split(" ").index(item)
        return tuple(it if i != idx else value for i, it in enumerate(batch))

    @staticmethod
    def _getitem_index(idx, _=None):
        return idx

    @staticmethod
    def _getitem_from_ctx(_, ctx, ctx_key):
        return ctx[ctx_key]

    def __getitem__(self, idx):
        if isinstance(idx, slice):
            return [self[i] for i in range(len(self))[idx]]
        if isinstance(idx, list):
            return [self[i] for i in idx]
        if idx < 0:
            idx = len(self) + idx

        items = []
        ctx = {} if self.propagate_ctx else None
        for getitem_fn in self._getitem_fns:
            item = getitem_fn(idx, ctx)
            items.append(item)
        if len(items) == 1:
            # single item -> no tuple
            items = items[0]
        else:
            # multiple items -> wrap into tuple
            items = tuple(items)
        if self.return_ctx:
            return items, ctx
        return items

    def __len__(self):
        return len(self.dataset)

    @property
    def root_dataset(self):
        # root_dataset is implemented in base class -> not handled in __getattr__
        return self.dataset.root_dataset

    def has_wrapper(self, wrapper):
        if self == wrapper:
            return True
        return self.dataset.has_wrapper(wrapper)

    def has_wrapper_type(self, wrapper_type):
        if type(self) == wrapper_type:
            return True
        return self.dataset.has_wrapper_type(wrapper_type)

    @property
    def all_wrappers(self):
        return [self] + self.dataset.all_wrappers

    @property
    def all_wrapper_types(self):
        return [type(self)] + self.dataset.all_wrapper_types

    def get_wrappers_of_type(self, wrapper_type):
        wrappers = self.dataset.get_wrappers_of_type(wrapper_type)
        if type(self) == wrapper_type:
            return [self] + wrappers
        return wrappers

    def worker_init_fn(self, rank, **kwargs):
        self.dataset.worker_init_fn(rank, **kwargs)

    def __getattr__(self, item):
        if item == "dataset":
            return getattr(super(), item)
        return getattr(self.dataset, item)

    def __iter__(self):
        """ torch.utils.data.Dataset doesn't define __iter__ which makes 'for sample in dataset' run endlessly """
        for i in range(len(self)):
            yield self[i]
