import numpy as np
from scipy.stats import norm


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

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

        self.upper_bound = upper_bound
        self.lower_bound = lower_bound

        self.user_model = None
        self.proposal_model_mean = np.array((nr_dims, 1))
        self.proposal_model_covariance = np.diag(np.ones((nr_dims, )) * 5)

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

    def initialize(self):
        pass

    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 = self.rng.multivariate_normal(
                self.proposal_model_mean, 
                self.proposal_model_covariance
            )

            # 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):
        pass

    def update_user_preference_model(self):
        pass

    def update_proposal_model(self):
        pass


if __name__ == '__main__':
    acquisition = PreferenceExpectedImprovement(10, 2, -1.0, 1.0)