#!/usr/bin/env python

import numpy as np
from olympus.surfaces import AbstractSurface
from scipy.stats import multivariate_normal
from olympus import Logger


class GaussianMixture(AbstractSurface):

    def __init__(self, param_dim=2, num_gauss=2, cov_scale=0.5, diagonal_cov=False, means=None, covariances=None,
                 noise=None, random_seed=None):
        """Random Gaussian Mixture surface. Given the number of dimensions and Gaussian distributions provided,
        the surface is the sum of probability densities. Unless a set of means and covariances are provided,
        a set of random means and covariances are generated.

        Args:
            param_dim (int): Number of input dimensions. Default is 2.
            num_gauss (int): Number of Gaussians. Default is 2.
            cov_scale (float): Scaling factor for the randomly generated covariance matrices (default is 0.5).
                A random covariance matrix B is generated by multiplying a random matrix A by its transpose.
                The elements of the random matrix are drawn from a uniform distribution. `cov_scale` is used to scale
                the resulting covariance matrix. Setting it to one typically generates smooth surfaces, while setting
                it to 0.1 generates rougher surfaces.
            diagonal_cov (bool): Whether to use a diagonal rather than full covariance matrix. Default is False.
            means (array): 2-dimensional array (shape=(#gaussians, #dimensions)) providing a list of means for the
                Gaussians. If provided, `param_dim` and `num_gauss` will be ignored, as they will be derived from the
                dimensionality of this array.
            covariances (array): 3-dimensional array (shape=(#gaussians, #dimensions, #dimensions)) providing a list
                of covariance matrices for the Gaussians. If provided, `param_dim` and `num_gauss` will be ignored,
                as they will be derived from the dimensionality of this array.
            noise (Noise): Noise object that injects noise into the evaluations of the surface. Default is None.
            random_seed (int): Fix the random seed for reproducible surfaces. Default is None, i.e. random.
        """
        param_dim, _num_gauss = self._determine_param_dim_and_num_gauss(param_dim, num_gauss, means, covariances)
        AbstractSurface.__init__(**locals())
        np.random.seed(self.random_seed)  # set random seed
        self._create_gaussian_mixture()

    @staticmethod
    def _determine_param_dim_and_num_gauss(param_dim, num_gauss, means, covariances):

        # no means and covs provided
        if means is None and covariances is None:
            _num_gauss = num_gauss
            _param_dim = param_dim
        # means provided
        elif means is not None and covariances is None:
            _num_gauss = np.shape(means)[0]
            _param_dim = np.shape(means)[1]
        # covs provided
        elif means is not None and covariances is not None:
            _num_gauss = np.shape(covariances)[0]
            _param_dim = np.shape(covariances)[1]
        # both means and covs provided
        elif means is not None and covariances is not None:
            assert np.shape(means)[0] == np.shape(covariances)[0]
            assert np.shape(means)[1] == np.shape(covariances)[1]
            _num_gauss = np.shape(means)[0]
            _param_dim = np.shape(means)[1]

        return _param_dim, _num_gauss

    @property
    def minima(self):
        # needs to be found by optimization
        message = 'Unknown minima: these need to be found numerically'
        Logger.log(message, 'WARNING')
        return None

    @property
    def maxima(self):
        # needs to be found by optimization
        message = 'Unknown maxima: these need to be found numerically'
        Logger.log(message, 'WARNING')
        return None

    def _run(self, params):
        value = - np.sum([g.pdf(params) for g in self.gaussians], axis=0)
        if self.noise is None:
            return value
        else:
            return self.noise(value)

    def _create_gaussian_mixture(self):

        # -------------------------
        # Get means and covariances
        # -------------------------
        # if means are provided
        if self.means is not None:
            self._means = np.array(self.means)
        # else generate random ones
        else:
            self._generate_random_means()

        # if covariances provided
        if self.covariances is not None:
            self._covs = np.array(self.covariances)
        # else generate random ones
        else:
            self._generate_random_covs()

        # ----------------
        # Create Gaussians
        # ----------------
        self.gaussians = []
        for i in range(self._num_gauss):
            g = multivariate_normal(mean=self._means[i], cov=self._covs[i])
            self.gaussians.append(g)

    def _generate_random_means(self):
        self._means = []
        for i in range(self._num_gauss):
            mean = np.random.uniform(low=0, high=1, size=len(self.param_space))
            self._means.append(mean)

    def _generate_random_covs(self):
        """

        Args:
            size:
            scale (float): scale the random covariance matrix.

        Returns:

        """
        self._covs = []

        if self.diagonal_cov is False:
            for i in range(self._num_gauss):
                # generate random positive semi-definite matrix, scale by sigma_scale
                A = np.random.rand(len(self.param_space), len(self.param_space))
                alpha = 0.01
                cov = np.dot(A, A.transpose()) + alpha * np.diag(np.ones(len(self.param_space)))
                self._covs.append(cov * self.cov_scale)
        elif self.diagonal_cov is True:
            for i in range(self._num_gauss):
                diag = np.random.uniform(low=0, high=1, size=len(self.param_space)) * self.cov_scale
                cov = np.diag(np.array(diag))
                self._covs.append(cov)

