# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/environments/20_data_generators.ipynb.

# %% auto 0
__all__ = ['make_binary', 'make_regular_binary']

# %% ../../nbs/environments/20_data_generators.ipynb 4
# General libraries:
import numpy as np

# %% ../../nbs/environments/20_data_generators.ipynb 5
def make_binary(n,
                num_features,
                num_informative_features=None,
                proportions = None,
                coef_range = (-1, 1),
                return_Y = False,
                seed = None
                ):

    """
    Generates a dataset with binary features. The first num_informative_features features are informative, the rest are noise.

    Args:
        n (int): Number of samples
        num_features (int): Number of features
        num_informative_features (int, optional): Number of informative features. Defaults to None.
        proportions (list, optional): Proportions of 0s and 1s for each feature. Defaults to None. If None, the proportions are drawn from a uniform distribution.
        return_Y (bool, optional): Weather to generate and return labels in addition to the features and coefficients. Defaults to False.
        seed (int, optional): Random seed. Defaults to None.

        
    Returns:
        coefficients (np.array): The coefficients of the features
        X (np.array): The generated features
    """


    if seed != None:
        np.random.seed(seed)

    if num_informative_features == None:
        num_informative_features = num_features

    # Generate X:
    if proportions == None:
            prop_0 = np.random.uniform(0,1, size = (num_features))
            prop_1 = 1 - prop_0
    else:
        prop_0 = proportions
        prop_1 = 1 - proportions
    
    X = np.zeros((n, num_features))
    for i in range(num_features):
        X[:,i] = np.random.choice([0,1], size = (n), p = [prop_0[i], prop_1[i]])

    # draw num_informative_features features 
    informative_features = np.random.choice(range(num_features), size = (num_informative_features), replace = False)

    coefficients = np.zeros((num_features))
    coefficients[informative_features] = np.random.uniform(coef_range[0], coef_range[1], size = (num_informative_features))

    if return_Y:
        return coefficients, X, Y
    else:
        return coefficients, X

# %% ../../nbs/environments/20_data_generators.ipynb 6
def make_regular_binary(n,
                        num_0 = 5,
                        num_1 = 2,
                        coef_range = (-1, 1),
                        seed = None
                        ):

    """
    Generate binary features with repeating pattern (e.g., to simulate weekdays and weekends).
    It will always start with 0 for num_0 times followed by 1 for num_1 times.

    Args:
        n (int): Number of samples
        num_0 (int, optional): Number of 0s in the pattern. Defaults to 5.
        num_1 (int, optional): Number of 1s in the pattern. Defaults to 2.
        coef_range (tuple, optional): Range of the coefficients. Defaults to (-1, 1).
        seed (int, optional): Random seed. Defaults to None.

    Returns:
        coefficients (np.array): The coefficients of the features
        X (np.array): The generated features

    """

    if seed != None:
        np.random.seed(seed)

    X = np.zeros((n))
    for i in range(n):
        if i % (num_0 + num_1) < num_0:
            X[i] = 0
        else:
            X[i] = 1
        
    coefficients = np.zeros((1))
    coefficients[0] = np.random.uniform(coef_range[0], coef_range[1], size = (1))

    return coefficients, X
