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


class WineDataset(ClassicDataset):

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

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

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