########## Packages ##########
import numpy as np
import plotly.graph_objects as go 
from math import e 
##############################

class Test_function():
    """ 
    The class of test functions
    """
    def rastrigin(self, x):
        """
        Rastrigin function
        """
        val = 10.0*x.size 
        for i in range(x.size):
            val += x[i]**2 - 10*np.cos(2*np.pi*x[i])
        return val 

    def ackley(self, x):
        """
        Ackley function
        """
        val = -20*np.exp(-0.2*(np.sqrt(0.5*(x[0]**2 + x[1]**2))))
        val += -np.exp(0.5*(np.cos(2*np.pi*x[0]) + np.cos(2*np.pi*x[1])))
        val += e + 20 
        return val 

class Plot_func():
    """
    The class to plot functions
    """
    def func_value(self, func, x_range, y_range):
        """
        get the value of the function (2 dim)
        Input:
            func: function
            x_range, y_range: numpy linspace
        Output:
            X, Y: numpy meshgrid
            Z: corresponding func value
        """
        X, Y = np.meshgrid(x_range, y_range)
        Z = np.empty_like(X)
        for i in range(len(X)):
            for j in range(len(X[i])):
                vec_x = np.array([X[i][j], Y[i][j]])
                Z[i][j] = func(vec_x)
        return X, Y, Z
    
    def plotly_surface_contour(self, func, x_range, y_range):
        """
        plot the surface of the function using plotly,
        https://plotly.com/python/3d-surface-plots/
        """
        ##### get the value of the function #####
        X, Y, Z = self.func_value(func, x_range, y_range)
        #########################################

        ##### plot #####
        fig = go.Figure(data=[go.Surface(z=Z, x=X, y=Y)])
        fig.update_traces(contours_z=dict(show=True, usecolormap=True,
                            highlightcolor="limegreen", project_z=True))
        fig.update_layout(autosize=False,
                            scene_camera_eye=dict(x=1.87, y=0.88, z=-0.64),
                            width=500, height=500,
                            margin=dict(l=65, r=50, b=65, t=90))
        fig.show()
        ################

##### instance of the class #####
test = Test_function()
pf = Plot_func()
#################################

##### Plot Rastrigin function #####
x = np.linspace(-5.12, 5.12, 100)
y = np.linspace(-5.12, 5.12, 100)
pf.plotly_surface_contour(
    func=test.rastrigin, 
    x_range=x,
    y_range=y)
###################################

##### Plot Ackley function #####
x = np.linspace(-5, 5, 100)
y = np.linspace(-5, 5, 100)
pf.plotly_surface_contour(
    func=test.ackley, 
    x_range=x, 
    y_range=y)
################################