# optimal score ~ 0.67366

import matplotlib.pyplot as plt
import matplotlib.cm as cm

import numpy as np
import math
import random
import time


# from: http://extremelearning.com.au/how-to-generate-uniformly-random-points-on-n-spheres-and-n-balls/
def muller(N, d):
    points = []
    while len(points) < N:
        u = np.random.normal(0, 1, d)
        norm = np.sum(u ** 2) ** (0.5)
        r = random.random() ** (1.0 / d)
        pt = r * u / norm
        points.append(tuple(pt))
    return points

def tweak(point, delta):
    change = muller(1, 2)[0]
    new_x = point[0] + delta * change[0]
    new_y = point[1] + delta * change[1]
    return (new_x, new_y)


plt.ion()
plt.style.use('ggplot')
fig, (ax_contour, ax_score) = plt.subplots(1,2)

score_line, = ax_score.plot([], [])
xdata = []
ydata = []

explore_line, = ax_contour.plot([], [], linestyle='-', marker='.', color='red')
xline = []
yline = []
best = None


plt.subplots_adjust(bottom=0.15)
ax_contour.set_title("Contour Plot", fontsize=20)
ax_score.set_title("Value", fontsize=20)
plt.show(block=False)

props = dict(boxstyle="round", facecolor="blue", alpha=0.5)
text = ax_score.text(
    0,
    -0.1,
    f"Best Score: {best}",
    transform=ax_score.transAxes,
    fontsize=14,
    verticalalignment="top",
    bbox=props,
)

delta = 0.01
x = np.arange(-10, 10, delta)
y = np.arange(-10, 10, delta)
X, Y = np.meshgrid(x, y)
Z = np.sin(3*np.pi*X)**2 + (X-1)**2*(1+np.sin(3*np.pi*Y)**2) + (Y-1)**2*(1+np.sin(2*np.pi*Y)**2)
new_line = True
levels = np.arange(0, 100, 5)


CS = ax_contour.contour(X, Y, Z, levels=levels)

plt.pause(0.0001)

pt_x = random.random() * 20 - 10
pt_y = random.random() * 20 - 10
value = math.sin(3*math.pi*pt_x)**2 + (pt_x-1)**2*(1+math.sin(3*math.pi*pt_y)**2) + (pt_y-1)**2*(1+math.sin(2*math.pi*pt_y)**2)
loops = 0

while True:
    
    if new_line:
        if loops > 20:
            plt.show(block=True)
        loops += 1
        explore_line, = ax_contour.plot([], [], linestyle='-', marker='.', color='red')
        xline = [pt_x]
        yline = [pt_y]
        explore_line.set_data(xline, yline)
        new_line = False

    else:
        xline.append(pt_x)
        yline.append(pt_y)
        explore_line.set_data(xline, yline)
    
    if best is None or value < best:
        best = value
        text.remove()
        text = ax_score.text(
            0,
            -0.05,
            f"Best Score: {best}",
            transform=ax_score.transAxes,
            fontsize=14,
            verticalalignment="top",
            bbox=props,
        )

    xdata.append(len(xdata))
    ydata.append(value)
    score_line.set_data(xdata, ydata)
    cur_x = ax_score.get_xlim()
    cur_y = ax_score.get_ylim()

    if xdata[-1] > cur_x[1]:
        ax_score.set_xlim((0, cur_x[1] + 100))
    if ydata[-1] > cur_y[1]:
        ax_score.set_ylim((0, max(ydata)))

    

    # time.sleep()
    plt.pause(0.00001)# ax_contour.clabel(CS, inline=True, fontsize=10)

    new_x = pt_x
    new_y = pt_y
    new_value = value

    tries = 0
    while new_value >= value:

        tries += 1
        if tries > 10000:
            plt.pause(0.5)
            tries = 0
            new_x = random.random() * 12 - 6
            new_y = random.random() * 12 - 6
            new_value = new_x ** 2 + new_y ** 2 + 25 * (math.sin(new_x)**2 + math.sin(new_y)**2)
            new_line = True
            break


        new_point = tweak((pt_x, pt_y), 0.1)
        new_x = new_point[0]
        new_y = new_point[1]
        new_value = math.sin(3*math.pi*new_x)**2 + (new_x-1)**2*(1+math.sin(3*math.pi*new_y)**2) + (new_y-1)**2*(1+math.sin(2*math.pi*new_y)**2)

    
    # print(value, '->', new_value)
    # print((pt_x, pt_y), '->', (new_x, new_y))

    pt_x, pt_y = new_x, new_y
    value = new_value