import numpy as np
from scipy.stats import norm, multivariate_normal


class PreferenceExpectedImprovement:
    def __init__(self, nr_dims, nr_samples, lower_bound, upper_bound, initial_variance, update_variance, seed=None):

        self.nr_dims = nr_dims
        self.nr_samples = int(nr_samples)

        self.lower_bound = lower_bound
        self.upper_bound = upper_bound

        self.init_var = initial_variance
        self.update_var = update_variance

        self.rng = np.random.default_rng(seed=seed)

        # initial proposal distribution
        self.proposal_mean = np.zeros((nr_dims, 1))
        self.proposal_cov = np.diag(np.ones((nr_dims,)) * self.init_var)

    def rejection_sampling(self):
        samples = np.empty((0, self.nr_dims))
        while samples.shape[0] < self.nr_samples:
            # sample from the multi variate gaussian distribution
            sample = np.zeros((1, self.nr_dims))
            for i in range(self.nr_dims):
                check = False
                while not check:
                    sample[0, i] = self.rng.normal(self.proposal_mean[i], self.proposal_cov[i, i])
                    if self.lower_bound <= sample[0, i] <= self.upper_bound:
                        check = True

            samples = np.append(samples, sample, axis=0)

        return samples

    def expected_improvement(self, gp , X, kappa=0.01):
        X_sample = self.rejection_sampling()

        mu_sample, sigma_sample = gp.predict(X_sample, return_std=True)
        sigma_sample = sigma_sample.reshape(-1, 1)

        mu = gp.predict(X)
        mu_best = np.max(mu)

        with np.errstate(divide='warn'):
            imp = mu_sample - mu_best - kappa
            imp = imp.reshape(-1, 1)
            z = imp / sigma_sample
            ei = imp * norm.cdf(z) + sigma_sample * norm.pdf(z)
            ei[sigma_sample == 0.0] = 0.0

        idx = np.argmax(ei)
        x_next = X_sample[idx, :]

        return x_next

    def update_proposal_model(self, preference_mean, preference_bool):
        cov_diag = np.ones((self.nr_dims,)) * self.init_var
        cov_diag[preference_bool] = self.update_var

        preference_cov = np.diag(cov_diag)

        preference_mean = preference_mean.reshape(-1, 1)

        posterior_mean = np.linalg.inv(np.linalg.inv(self.proposal_cov) + np.linalg.inv(preference_cov))\
            .dot(np.linalg.inv(self.proposal_cov).dot(self.proposal_mean)
                 + np.linalg.inv(preference_cov).dot(preference_mean))

        posterior_cov = np.linalg.inv(np.linalg.inv(self.proposal_cov) + np.linalg.inv(preference_cov))

        self.proposal_mean = posterior_mean
        self.proposal_cov = posterior_cov