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

from anti_kd_backdoor.data.dataset.base import (
    DatasetInterface,
    IndexDataset,
    IndexRatioDataset,
    PoisonLabelDataset,
    RangeRatioDataset,
    RangeRatioPoisonLabelDataset,
    RatioDataset,
    RatioPoisonLabelDataset,
)
from anti_kd_backdoor.data.dataset.types import XY_TYPE


class FakeDataset(DatasetInterface, Dataset):
    cache: dict[tuple, tuple[torch.Tensor, list[int]]] = dict()

    def __init__(self,
                 *,
                 x_shape: tuple[int, int, int] = (3, 32, 32),
                 y_range: tuple[int, int] = (0, 9),
                 nums: int = 10000,
                 cache_xy: bool = False) -> None:
        self._nums = nums
        self._x_shape = x_shape
        self._y_range = y_range
        self._raw_num_classes = y_range[1] - y_range[0] + 1

        if cache_xy:
            cache_key = (x_shape, y_range, nums)
            if cache_key not in FakeDataset.cache:
                x, y = self._prepare_xy()
                FakeDataset.cache[cache_key] = (x, y)
            x, y = FakeDataset.cache[cache_key]
        else:
            x, y = self._prepare_xy()

        self.data, self.targets = x, y

    def __len__(self) -> int:
        return len(self.targets)

    def __getitem__(self, index: int) -> tuple[torch.Tensor, torch.Tensor]:
        x = self.data[index]
        y = self.targets[index]

        return x, y

    def get_xy(self) -> XY_TYPE:
        return list(self.data), self.targets.copy()

    def set_xy(self, xy: XY_TYPE) -> None:
        x, y = xy
        assert len(x) == len(y)

        self.data = np.stack(x, axis=0)
        self.targets = y.copy()

    @property
    def num_classes(self) -> int:
        return len(set(self.targets))

    @property
    def raw_num_classes(self) -> int:
        return self._raw_num_classes

    def _prepare_xy(self) -> tuple[torch.Tensor, list[int]]:
        x = torch.rand((self._nums, *self._x_shape))
        num_per_class = self._nums // self._raw_num_classes
        y = [
            i for _ in range(num_per_class)
            for i in range(self._raw_num_classes)
        ]
        y.extend([self._y_range[0]] * (self._nums - len(y)))

        return x, y


class IndexFakeDataset(FakeDataset, IndexDataset):

    def __init__(self, *, start_idx: int, end_idx: int, **kwargs) -> None:
        FakeDataset.__init__(self, **kwargs)
        IndexDataset.__init__(self, start_idx=start_idx, end_idx=end_idx)


class RatioFakeDataset(FakeDataset, RatioDataset):

    def __init__(self, *, ratio: float, **kwargs) -> None:
        FakeDataset.__init__(self, **kwargs)
        RatioDataset.__init__(self, ratio=ratio)


class RangeRatioFakeDataset(FakeDataset, RangeRatioDataset):

    def __init__(self, *, range_ratio: tuple[float, float], **kwargs) -> None:
        FakeDataset.__init__(self, **kwargs)
        RangeRatioDataset.__init__(self, range_ratio=range_ratio)


class IndexRatioFakeDataset(FakeDataset, IndexRatioDataset):

    def __init__(self, *, start_idx: int, end_idx: int, ratio: float,
                 **kwargs) -> None:
        FakeDataset.__init__(self, **kwargs)
        IndexRatioDataset.__init__(self,
                                   start_idx=start_idx,
                                   end_idx=end_idx,
                                   ratio=ratio)


class PoisonLabelFakeDataset(FakeDataset, PoisonLabelDataset):

    def __init__(self, *, poison_label: int, **kwargs) -> None:
        FakeDataset.__init__(self, **kwargs)
        PoisonLabelDataset.__init__(self, poison_label=poison_label)


class RatioPoisonLabelFakeDataset(FakeDataset, RatioPoisonLabelDataset):

    def __init__(self, *, ratio: float, poison_label: int, **kwargs) -> None:
        FakeDataset.__init__(self, **kwargs)
        RatioPoisonLabelDataset.__init__(self,
                                         ratio=ratio,
                                         poison_label=poison_label)


class RangeRatioPoisonLabelFakeDataset(FakeDataset,
                                       RangeRatioPoisonLabelDataset):

    def __init__(self, *, range_ratio: tuple[float, float], poison_label: int,
                 **kwargs) -> None:
        FakeDataset.__init__(self, **kwargs)
        RangeRatioPoisonLabelDataset.__init__(self,
                                              range_ratio=range_ratio,
                                              poison_label=poison_label)


FAKE_DATASETS_MAPPING = {
    'FakeDataset': FakeDataset,
    'IndexFakeDataset': IndexFakeDataset,
    'RatioFakeDataset': RatioFakeDataset,
    'RangeRatioFakeDataset': RangeRatioFakeDataset,
    'IndexRatioFakeDataset': IndexRatioFakeDataset,
    'PoisonLabelFakeDataset': PoisonLabelFakeDataset,
    'RatioPoisonLabelFakeDataset': RatioPoisonLabelFakeDataset,
    'RangeRatioPoisonLabelFakeDataset': RangeRatioPoisonLabelFakeDataset
}


def build_fake_dataset(dataset_cfg: dict) -> FakeDataset:
    if 'type' not in dataset_cfg:
        raise ValueError('Dataset config must have `type` field')
    dataset_type = dataset_cfg.pop('type')
    if dataset_type not in FAKE_DATASETS_MAPPING:
        raise ValueError(
            f'Dataset `{dataset_type}` is not support, '
            f'available datasets: {list(FAKE_DATASETS_MAPPING.keys())}')
    dataset = FAKE_DATASETS_MAPPING[dataset_type]

    return dataset(**dataset_cfg)