from core.datasets.classic.ClassicDataset import ClassicDataset
from sklearn.datasets import load_iris


class IrisDataset(ClassicDataset):

    def __init__(self, seed):
        self.dataset = load_iris()

        super().__init__(
            name='iris',
            num_classes=3,
            num_features=4,
            sample_shape=(4,),
            feature_names=self.dataset.feature_names,
            seed=seed
        )

    def _get_x_y(self):
        x = self.dataset.data
        y = self.dataset.target
        return x, y
