import numpy as np
from scipy.stats import norm


def generate_samples(n_samples, use_gaussian, means, variances, lower, upper, seed):
    """
    Generate samples from a uniform or Gaussian distribution based on a boolean array.

    Args:
        n_samples: The number of samples to generate.
        use_gaussian: A boolean array of shape (n_dims,). If use_gaussian[i] is True,
                      generate the i-th dimension of the samples from a Gaussian distribution.
        means: An array of shape (n_dims,) giving the means of the Gaussian distributions.
        variances: An array of shape (n_dims,) giving the variances of the Gaussian distributions.
        lower: lower bound of the samples.
        upper: upper bound of the samples.
        seed: seed of the random generator.

    Returns:
        samples: An array of shape (n_samples, n_dims) containing the generated samples.
    """
    n_dims = len(use_gaussian)
    samples = np.empty((n_samples, n_dims))
    rng = np.random.default_rng(seed=seed)
    for i in range(n_dims):
        if use_gaussian[i]:
            samples[:, i] = rng.normal(means[i], np.sqrt(variances[i]), n_samples).clip(min=lower, max=upper)
        else:
            samples[:, i] = rng.uniform(lower, upper, n_samples)
    return samples


def expected_improvement(gp, X, n_samples, use_gaussian, means, variances, kappa=0.01, lower=-1.0, upper=1.0, seed=None):
    """
    Compute the expected improvement at points X based on existing samples X_sample
    using a Gaussian process surrogate model.

    Args:
        X: Points at which the objective function has already been sampled.
        n_samples: The number of samples to generate.
        gp: GaussianProcessRegressor object.
        use_gaussian: A boolean array of shape (n_dims,). If use_gaussian[i] is True,
                      generate the i-th dimension of the samples from a Gaussian distribution.
        means: An array of shape (n_dims,) giving the means of the Gaussian distributions.
        variances: An array of shape (n_dims,) giving the variances of the Gaussian distributions.
        kappa: Exploitation-exploration trade-off parameter.
        lower: lower bound of the samples.
        upper: upper bound of the samples.
        seed: seed of the random generator.

    Returns:
        next best observation
    """
    # Generate samples
    X_test = generate_samples(n_samples, use_gaussian, means, variances, lower, upper, seed)

    mu, sigma = gp.predict(X_test, return_std=True)
    mu_sample = gp.predict(X)

    sigma = sigma.reshape(-1, 1)

    # Needed for noise-based model,
    # otherwise use np.max(Y_sample).
    mu_sample_opt = np.max(mu_sample)

    with np.errstate(divide='warn'):
        imp = mu - mu_sample_opt - kappa
        imp = imp.reshape(-1, 1)
        Z = imp / sigma
        ei = imp * norm.cdf(Z) + sigma * norm.pdf(Z)
        ei[sigma == 0.0] = 0.0

    # Return the sample with the maximum expected improvement
    idx = np.argmax(ei)
    X_next = X_test[idx, :]

    return X_next