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

import numpy as np


class LeafDataset(ClassicDataset):
    def __init__(self, seed):
        self.dataset = fetch_openml(data_id=1482, as_frame=True)

        super().__init__(
            name='leaf',
            num_classes=30,
            num_features=15,
            sample_shape=(15,),
            feature_names=self.dataset.feature_names,
            seed=seed
        )

    def _get_x_y(self):
        x = np.asarray(self.dataset.data)
        y = np.asarray(self.dataset.target)

        # Parse the target to numbers.
        y = y.astype(int) - 1

        return x, y
