# encoding=utf8
import logging
from scipy.stats import norm
import numpy as np

from robo.acquisition.base import AcquisitionFunction

logger = logging.getLogger(__name__)

class EI(AcquisitionFunction):
    r"""
        Expected Improvement computes for a given x the acquisition value by
        :math:`EI(X) := \mathbb{E}\left[ \max\{0, f(\mathbf{X^+}) - f_{t+1}(\mathbf{X}) - \xi\right] \} ]`, with
        :math:`f(X^+)` as the incumbent.

        :param model: A model that implements at least

                 - predict(X)
                 - getCurrentBestX().

               If you want to calculate derivatives than it should also support

                 - predictive_gradients(X)
        :param X_lower: Lower bounds for the search, its shape should be 1xD (D = dimension of input space)
        :type X_lower: np.ndarray (1,D)
        :param X_upper: Upper bounds for the search, its shape should be 1xD (D = dimension of input space)
        :type X_upper: np.ndarray (1,D)
        :param compute_incumbent: A python function that takes as input a model and returns a np.array as incumbent
        :param par: Controls the balance between exploration and exploitation of the acquisition
                    function. Default is 0.01
    """

    long_name = "Expected Improvement"

    def __init__(self, model, X_lower, X_upper, compute_incumbent, par=0.01, **kwargs):
        #self.model = model
        self.par = par
        #self.X_lower = X_lower
        #self.X_upper = X_upper
        self.compute_incumbent = compute_incumbent
        super(EI, self).__init__(model, X_lower, X_upper)

        logger.debug("Test")

    def compute(self, X, derivative=False, **kwargs):
        """
        Computes the EI value and its derivatives.

        :param X: The point at which the function is to be evaluated.
        :type X: np.ndarray (1,D)
        :param derivative: This controls whether the derivative is to be returned.
        :type derivative: Boolean
        :return: The value of EI and optionally its derivative at X.
        :rtype: np.ndarray(1, 1) or (np.ndarray(1, 1), np.ndarray(1, D))
        """

        if X.shape[0] > 1:
            raise ValueError("EI is only for single test points")

        if len(X.shape) == 1:
            X = X[:, np.newaxis]

        if np.any(X < self.X_lower) or np.any(X > self.X_upper):
            if derivative:
                f = 0
                df = np.zeros((1, X.shape[1]))
                return np.array([[f]]), np.array([df])
            else:
                return np.array([[0]])

        m, v = self.model.predict(X, full_cov=True)

        # Use the best seen observation as incumbent
        incumbent, eta = self.compute_incumbent(self.model)

        s = np.sqrt(v)
        z = (eta - m - self.par) / s
        f = (eta - m - self.par) * norm.cdf(z) + s * norm.pdf(z)
        if derivative:
            dmdx, ds2dx = self.model.predictive_gradients(X)
            dmdx = dmdx[0]
            ds2dx = ds2dx[0][:, None]
            dsdx = ds2dx / (2 * s)
            df = (-dmdx * norm.cdf(z) + (dsdx * norm.pdf(z))).T
        if (f < 0).any():
            f[np.where(f < 0)] = 0.0
            if derivative:
                df[np.where(f < 0), :] = np.zeros_like(X)
        if (f < 0).any():
            raise Exception
        if len(f.shape) == 1:
            f = np.array([f])
        if derivative:
            if len(df.shape) == 3:
                return_df = df
            else:
                return_df = np.array([df])
            return f, return_df
        else:
            return f

@@ -2674,25 +2674,17 @@
         
-incumbent
+_
 , eta = 
@@ -2742,16 +2742,137 @@
 sqrt(v)%0A
+%0A        if (s == 0).any():%0A            f = np.array(%5B%5B0%5D%5D)%0A            df = np.zeros((1, X.shape%5B1%5D))%0A        else:%0A    
         
@@ -2900,16 +2900,20 @@
 ar) / s%0A
+    
         
@@ -2969,32 +2969,36 @@
 .pdf(z)%0A        
+    
 if derivative:%0A 
@@ -3004,24 +3004,28 @@
             
+   
+ 
 dmdx, ds2dx 
@@ -3053,32 +3053,36 @@
 ve_gradients(X)%0A
+    
             dmdx
@@ -3104,16 +3104,20 @@
         
+    
 ds2dx = 
@@ -3146,16 +3146,20 @@
         
+ 
+   
 dsdx = d
@@ -3173,16 +3173,20 @@
 (2 * s)%0A
+    
         
@@ -3241,32 +3241,36 @@
 (z))).T%0A        
+   
+ 
 if (f %3C 0).any()
@@ -3287,144 +3287,73 @@
     
-f%5Bnp.where(f %3C 0)%5D = 0.0%0A            if derivative:%0A                df%5Bnp.where(f %3C 0), :%5D = np.zeros_like(X)%0A        if (f %3C 0).any():%0A
+    logger.error(%22Expected Improvement is smaller than 0!%22)%0A     
     
@@ -3355,25 +3355,24 @@
             
- 
 raise Except
@@ -3371,24 +3371,28 @@
 e Exception%0A
+    
         if l
@@ -3413,32 +3413,36 @@
  1:%0A            
+    
 f = np.array(%5Bf%5D
@@ -3439,24 +3439,25 @@
 .array(%5Bf%5D)%0A
+%0A
         if d

return ei=0 if std is 0