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

import numpy as np


class DiabetesDataset(ClassicDataset):

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

        # Binarize target
        median = np.median(self.dataset.target)
        self.dataset.target = (self.dataset.target >= median).astype(int)

        super().__init__(
            name='diabetes',
            num_classes=2,
            num_features=10,
            sample_shape=(10,),
            feature_names=self.dataset.feature_names,
            seed=seed
        )

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