import numpy as np
from scipy.stats import norm, multivariate_normal
import matplotlib.pyplot as plt


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

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

        # check if upper_bound and lower_bound are numpy arrays of shape (nr_dims, 1) or (nr_dims,) or if they are float
        self.upper_bound = upper_bound
        self.lower_bound = lower_bound

        self.initial_variance = init_var

        self.proposal_mean = np.zeros((nr_dims, 1))
        self.proposal_cov = np.diag(np.ones((nr_dims,)) * self.initial_variance)

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

    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)
            # sample = self.rng.multivariate_normal(
            #     mean,
            #     self.proposal_cov
            # )
            #
            # # check if the sample is within the bounds
            # if np.all(sample >= self.lower_bound) and np.all(sample <= self.upper_bound):
            #     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 likelihood(self, preference):
        covariance_diag = np.ones((self.nr_dims,)) * self.initial_variance
        covariance_diag[preference] = 0.05

        covariance = np.diag(covariance_diag)

        return covariance

    def update_proposal_model(self, preference_mean, preference_bool):

        covariance_diag = np.ones((self.nr_dims,)) * self.initial_variance
        covariance_diag[preference_bool] = 0.05

        preference_cov = np.diag(covariance_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))

        print(posterior_mean, posterior_cov)

        self.proposal_mean = posterior_mean
        self.proposal_cov = posterior_cov

    def plot_2D(self, mean, cov):
        print(mean.shape, cov.shape)
        if mean.shape == (2, 1):
            mean = mean.squeeze()

        gaussian = multivariate_normal(mean, cov)
        x = np.random.uniform(self.lower_bound, self.upper_bound, (self.nr_likelihood_samples, self.nr_dims))

        pdf = gaussian.pdf(x)
        pdf = pdf / pdf.sum()

        plt.scatter(x[:, 0], x[:, 1], c=pdf, cmap='viridis')

        # Add a color bar
        plt.colorbar(label='PDF value')

        # Add labels
        plt.xlabel('x1')
        plt.ylabel('x2')

        # Show the plot
        plt.show()


if __name__ == '__main__':
    # acquisition = PreferenceExpectedImprovement(10, 10, 10e4, -1.0, 1.0, 5.0)
    # sample_res = acquisition.rejection_sampling()
    # print(f"finished: {sample_res}")
    acquisition = PreferenceExpectedImprovement(10, 2, 10e4, -1.0, 1.0, 10.0)
    mean_ = np.array([0.5, 0.23])
    preference_ = [False, True]
    likelihood_cov = acquisition.likelihood(preference_)

    acquisition.plot_2D(acquisition.proposal_mean, acquisition.proposal_cov)
    acquisition.plot_2D(mean_, likelihood_cov)

    acquisition.update_proposal_model(mean_, preference_)
    acquisition.plot_2D(acquisition.proposal_mean, acquisition.proposal_cov)

    mean2 = np.array([0.33, 0.24])
    preference2 = [True, False]
    likelihood_cov2 = acquisition.likelihood(preference2)

    acquisition.plot_2D(mean2, likelihood_cov2)

    acquisition.update_proposal_model(mean2, preference2)
    acquisition.plot_2D(acquisition.proposal_mean, acquisition.proposal_cov)

    mean2 = np.array([-0.66, -0.5])
    preference2 = [False, False]
    likelihood_cov2 = acquisition.likelihood(preference2)

    acquisition.plot_2D(mean2, likelihood_cov2)

    acquisition.update_proposal_model(mean2, preference2)
    acquisition.update_proposal_model(mean2, preference2)
    acquisition.update_proposal_model(mean2, preference2)
    acquisition.update_proposal_model(mean2, preference2)
    acquisition.update_proposal_model(mean2, preference2)
    acquisition.plot_2D(acquisition.proposal_mean, acquisition.proposal_cov)