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

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

# %% ../../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))
    X = X.reshape(-1,1)

    return coefficients, X

# %% ../../nbs/environments/20_data_generators.ipynb 9
def make_ARMA(series, p=5, q=5, range_start = 1, range_end = -0.5, return_coefficients = False):

    """
    Generates an ARMA(p,q) series from a given series of noise around a mean.
     
    For the AR coefficients, the first third are drawn from a uniform distribution between 0 and 1,
    the rest from a uniform distribution between -0.5 and 0.

    Args:
        series (np.array): The series around which the ARMA series is generated
        p (int): The number of AR coefficients
        q (int): The number of MA coefficients
        return_coefficients (bool, optional): Weather to return the coefficients. Defaults to False.
        
    Returns:
        series_ARMA (np.array): The generated ARMA series
    """

    residuals = series - np.mean(series)

    series_MA = np.zeros_like(residuals)
    series_ARMA = np.zeros_like(residuals)

    non_stationary=True
    counter = 0
    while non_stationary:
        phi = np.zeros(p)
        phi[:int(p/3)] = np.random.uniform(0, range_start, int(p/3))
        phi[int(p/3):] = np.random.uniform(range_end, 0, p-int(p/3))
        if np.sum(np.abs(phi)) < 1:
            #print("phi: ", phi, "with sum: ", np.sum(phi), "after ", counter, "iterations")
            non_stationary=False
        if counter == 1_000_000:
            assert False, "phi not stationary"
        counter += 1

    theta = np.random.uniform(0, 1, q)

    for i in range(max(p, q), len(series_ARMA)):
        series_MA[i] = np.dot(residuals[i-q:i], theta)
        series_ARMA[i] = np.dot(series_ARMA[i-p:i], phi) + series_MA[i]
    
    series_ARMA = series_ARMA + np.mean(series)
    
    return series_ARMA

# %% ../../nbs/environments/20_data_generators.ipynb 11
def make_ARMA_features(num_samples, p=10, q=10, coef_range=(-1,1), data_mu = 0, data_sigma = 1):
    noise = np.random.normal(data_mu, data_sigma, num_samples)
    series = make_ARMA(noise, p=p, q=q)
    coefficients = np.random.uniform(coef_range[0], coef_range[1], size = (1))
    series = series.reshape(-1,1)

    return coefficients, series

# %% ../../nbs/environments/20_data_generators.ipynb 13
def make_continuous_iid(num_samples, coef_range=(-1,1), data_mu = 0, data_sigma = 1):
    series = np.random.normal(data_mu, data_sigma, num_samples)
    coefficient = np.random.uniform(coef_range[0], coef_range[1], size = (1))
    series = series.reshape(-1, 1)

    return coefficient, series
