from core.datasets.custom.CustomDataset import CustomDataset
from sklearn.datasets import make_classification


class CustomNDataset(CustomDataset):

    def __init__(self, num_features, seed):
        self.num_samples = 10000

        super().__init__(
            name=f'custom-{num_features}',
            num_classes=2,
            num_features=num_features,
            sample_shape=(num_features,),
            seed=seed
        )

    def _get_x_y(self):
        x, y = make_classification(
            n_samples=self.num_samples,
            n_features=self.num_features,
            n_informative=self.num_features // 2,
            n_redundant=0,
            n_repeated=0,
            n_classes=self.num_classes,
            random_state=self.seed
        )
        return x, y
