# Init matplotlib

%matplotlib inline

import matplotlib
matplotlib.rcParams['figure.figsize'] = (8, 8)

from mpl_toolkits.mplot3d import axes3d
import matplotlib.colors as colors

import numpy as np
import warnings

from scipy import optimize

"""
Explanation: Unconstrained global optimization with Scipy
Import required modules
End of explanation
"""


def plot_contour_2d_solution_space(func,
                                   fig=None,
                                   ax=None,
                                   show=True,
                                   xmin=-np.ones(2),
                                   xmax=np.ones(2),
                                   xstar=None,
                                   xvisited=None,
                                   title=""):
    """Plot points visited during the execution of an optimization algorithm.

    TODO
    """
    if (fig is None) or (ax is None):                # TODO
        fig, ax = plt.subplots(figsize=(12, 8))

    if xvisited is not None:
        xmin = np.amin(np.hstack([xmin.reshape([-1, 1]), xvisited]), axis=1)
        xmax = np.amax(np.hstack([xmax.reshape([-1, 1]), xvisited]), axis=1)

    x1_space = np.linspace(xmin[0], xmax[0], 200)
    x2_space = np.linspace(xmin[1], xmax[1], 200)

    x1_mesh, x2_mesh = np.meshgrid(x1_space, x2_space)

    zz = func(np.array([x1_mesh.ravel(), x2_mesh.ravel()])).reshape(x1_mesh.shape)

    ############################

    min_value = func(xstar)
    max_value = zz.max()

    levels = np.logspace(0.1, 3., 5)          # TODO

    im = ax.pcolormesh(x1_mesh, x2_mesh, zz,
                       vmin=0.1,              # TODO
                       vmax=max_value,
                       norm=colors.LogNorm(), # TODO
                       shading='gouraud',
                       cmap='gnuplot2') # 'jet' # 'gnuplot2'

    plt.colorbar(im, ax=ax)

    cs = plt.contour(x1_mesh, x2_mesh, zz, levels,
                     linewidths=(2, 2, 2, 2, 3),
                     linestyles=('dotted', '-.', 'dashed', 'solid', 'solid'),
                     alpha=0.5,
                     colors='white')
    ax.clabel(cs, inline=False, fontsize=12)

    ############################

    if xvisited is not None:
        ax.plot(xvisited[0],
                xvisited[1],
                '-og',
                alpha=0.5,
                label="$visited$")

    ############################

    if xstar is not None:
        sc = ax.scatter(xstar[0],
                   xstar[1],
                   c='red',
                   label="$x^*$")
        sc.set_zorder(10)        # put this point above every thing else

    ############################

    ax.set_title(title)

    ax.set_xlabel(r"$x_1$")
    ax.set_ylabel(r"$x_2$")

    ax.legend(fontsize=12)

    if show:
        plt.show()

    return fig, ax


def plot_2d_solution_space(func,
                           fig=None,
                           ax=None,
                           show=True,
                           xmin=-np.ones(2),
                           xmax=np.ones(2),
                           xstar=None,
                           xvisited=None,
                           angle_view=None,
                           title=""):
    """Plot points visited during the execution of an optimization algorithm.

    TODO
    """
    if fig is None or ax is None:                # TODO
        fig = plt.figure(figsize=(12, 8))
        ax = axes3d.Axes3D(fig)

    if angle_view is not None:
        ax.view_init(angle_view[0], angle_view[1])

    x1_space = np.linspace(xmin[0], xmax[0], 100)
    x2_space = np.linspace(xmin[1], xmax[1], 100)

    x1_mesh, x2_mesh = np.meshgrid(x1_space, x2_space)

    zz = func(np.array([x1_mesh.ravel(), x2_mesh.ravel()])).reshape(x1_mesh.shape)   # TODO

    ############################

    surf = ax.plot_surface(x1_mesh,
                           x2_mesh,
                           zz,
                           cmap='gnuplot2', # 'jet' # 'gnuplot2'
                           norm=colors.LogNorm(),   # TODO
                           rstride=1,
                           cstride=1,
                           #color='b',
                           shade=False)

    ax.set_zlabel(r"$f(x_1, x_2)$")

    fig.colorbar(surf, shrink=0.5, aspect=5)

    ############################

    if xstar is not None:
        ax.scatter(xstar[0],
                   xstar[1],
                   func(xstar),
                   #s=50,          # TODO
                   c='red',
                   alpha=1,
                   label="$x^*$")

    ax.set_title(title)

    ax.set_xlabel(r"$x_1$")
    ax.set_ylabel(r"$x_2$")

    ax.legend(fontsize=12)

    if show:
        plt.show()

    return fig, ax

"""
Explanation: Plot functions
End of explanation
"""


class Sphere:
    def __init__(self):
        self.reset_log()
        self.log = False
        self.arg_min = np.zeros(2)  # The actual optimal solution
    
    def reset_log(self):
        self.log_dict = {'x': [], 'y': []}
        
    def __call__(self, x):
        r"""The Sphere function.

        The Sphere function is a famous **convex** function used to test the performance of optimization algorithms.
        This function is very easy to optimize and can be used as a first test to check an optimization algorithm.

        $$
        f(\boldsymbol{x}) = \sum_{i=1}^{n} x_{i}^2
        $$

        Global minimum:
        $$
        f(\boldsymbol{0}) = 0
        $$

        Search domain:
        $$
        \boldsymbol{x} \in \mathbb{R}^n
        $$

        Example: single 2D point
        ------------------------

        To evaluate $x = \begin{pmatrix} 0 \\ 0 \end{pmatrix}$:

        >>> sphere( np.array([0, 0]) )
        0.0

        The result should be $f(x) = 0$.

        Example: single 3D point
        ------------------------

        To evaluate $x = \begin{pmatrix} 1 \\ 1 \\ 1 \end{pmatrix}$:

        >>> sphere( np.array([1, 1, 1]) )
        3.0

        The result should be $f(x) = 3.0$.

        Example: multiple 2D points
        ---------------------------

        To evaluate $x_1 = \begin{pmatrix} 0 \\ 0 \end{pmatrix}$,
        $x_2 = \begin{pmatrix} 1 \\ 1 \end{pmatrix}$ and
        $x_3 = \begin{pmatrix} 2 \\ 2 \end{pmatrix}$ at once:

        >>> sphere( np.array([[0, 1, 2], [0, 1, 2]]) )
        array([   0.,    2.,  8.])

        The result should be $f(x_1) = 0$, $f(x_2) = 1$ and $f(x_3) = 8$.

        Parameters
        ----------
        x : array_like
            One dimension Numpy array of the point at which the Sphere function is to be computed
            or a two dimension Numpy array of points at which the Sphere function is to be computed.

        Returns
        -------
        float or array_like
            The value(s) of the Sphere function for the given point(s) `x`.
        """
        # Remark: `sum(x**2.0)` is equivalent to `np.sum(x**2.0, axis=0)` but only the latter works if x is a scallar (e.g. x = np.float(3)).
        y = np.sum(x**2.0, axis=0)
        if self.log is True:
            self.log_dict['x'].append(x)
            self.log_dict['y'].append(y)
        return y

func = Sphere()    

xmin = np.full(2, -10)
xmax = np.full(2, 10)

"""
Explanation: Objective function
End of explanation
"""


from scipy import optimize

bounds = [[-10, 10], [-10, 10]]

res = optimize.differential_evolution(func,
                                      bounds,              # The initial point
                                      maxiter=100,         # The number of DE iterations
                                      polish=True)

print("x* =", res.x)
print("f(x*) =", res.fun)
print("Cause of the termination:", res.message)
print("Number of evaluations of the objective functions:", res.nfev)
print("Number of iterations performed by the optimizer:", res.nit)

print(res)

"""
Explanation: The "Differential Evolution" (DE) algorithm
Differential Evolution is a stochastic algorithm which attempts to find the global minimum of a function.
Official documentation:
* https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.differential_evolution.html#scipy.optimize.differential_evolution
More information:
* Practical advice
* Wikipedia article
Basic usage
End of explanation
"""


%%time

bounds = np.array([[-10, 10], [-10, 10]])

it_x_list = []
it_fx_list = []

def callback(xk, convergence):
    it_x_list.append(xk)
    it_fx_list.append(func(xk))
    print(len(it_x_list), xk, it_fx_list[-1])

func.reset_log()
func.log = True

with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    res = optimize.differential_evolution(func,
                                          bounds,              # The initial point
                                          maxiter=100,         # The number of DE iterations
                                          callback=callback,
                                          polish=False,
                                          disp=False)          # Print status messages

func.log = False

eval_x_array = np.array(func.log_dict['x']).T
eval_error_array = np.array(func.log_dict['y']) - func(func.arg_min)

it_x_array = np.array(it_x_list).T
it_error_array = np.array(it_fx_list) - func(func.arg_min)

print("x* =", res.x)
print("f(x*) =", res.fun)
print("Cause of the termination:", res.message)
print("Number of evaluations of the objective functions:", res.nfev)
print("Number of iterations performed by the optimizer:", res.nit)

plot_contour_2d_solution_space(func,
                               xmin=xmin,
                               xmax=xmax,
                               xstar=res.x,
                               xvisited=it_x_array,
                               title="Differential Evolution (best iterations solution only)");

plot_contour_2d_solution_space(func,
                               xmin=xmin,
                               xmax=xmax,
                               xstar=res.x,
                               xvisited=eval_x_array,
                               title="Differential Evolution (all tested solution)");

plt.loglog(it_error_array)
plt.title("Error of the best individual per iteration")
plt.xlabel("Iteration")
plt.ylabel("Error");

plt.loglog(eval_error_array)
plt.title("Error for all tested solutions")
plt.xlabel("Evaluation")
plt.ylabel("Error");

"""
Explanation: Performances analysis
End of explanation
"""


