import subprocess
from random import seed
import numpy as np
from mpl_toolkits.mplot3d import *
import matplotlib.pyplot as plt
from matplotlib import animation
from scipy.optimize import curve_fit


########################################################################
# General parameters
########################################################################

outfile = "test"
tmpfile = outfile+"_{0:03d}"

# xmin = -5.12
# xmax =  5.12
xmin = -3
xmax =  5
zmin = 0
zmax = 150
hres = 50
lres = 15
cnb = 20
samp_big_n = 100
gap = 50
ceil = zmax
view_alt = 32
view_angle = 105
nb_frames = 100
#nb_frames = 10

print(nb_frames,"frames")

np.random.seed(0)

def init():
    ########################################################################
    # Plot initialization
    ########################################################################

    fig = plt.figure(facecolor='black')
    ax = fig.gca(projection='3d', axisbg="black")
    ax.view_init(view_alt, view_angle)
    # ax.axis('off')

    # Black grid background
    ax.w_xaxis.set_pane_color((0,0,0))
    ax.w_yaxis.set_pane_color((0,0,0))
    ax.w_zaxis.set_pane_color((0,0,0))

    # Transparent axis grid lines
    ax.w_xaxis._axinfo.update({'grid' : {'color': (1,1,1, 0.2)}})
    ax.w_yaxis._axinfo.update({'grid' : {'color': (1,1,1, 0.2)}})
    ax.w_zaxis._axinfo.update({'grid' : {'color': (1,1,1, 0.2)}})

    # Transparent background:
    #fig.patch.set_alpha(0.)
    #ax.xaxis.set_visible(False)
    #ax.yaxis.set_visible(False)
    #ax.set_frame_on(False)

    plt.hold(True)

    return fig,ax


########################################################################
# Data grids and samples
########################################################################

# Low res grid
glx = np.arange(xmin, xmax, (xmax-xmin)/lres)
gly = np.arange(xmin, xmax, (xmax-xmin)/lres)
grid_low = np.meshgrid(glx, gly)

# Big res grid
ghx = np.arange(xmin, xmax, (xmax-xmin)/hres)
ghy = np.arange(xmin, xmax, (xmax-xmin)/hres)
grid_big = np.meshgrid(ghx, ghy)

# Big res sample
samp_big_x = np.random.random(samp_big_n)*(xmax-xmin)+xmin
samp_big_y = np.random.random(samp_big_n)*(xmax-xmin)+xmin
samp_big = (samp_big_x, samp_big_y)


########################################################################
# Fitness values
########################################################################

# "Real" problem: Rastrigin function
def rastrigin(x,y):
    a=10
    samp_big_n=2
    return a*samp_big_n + (x*x-a*np.cos(2*np.pi*x)) + (y*y-a*np.cos(2*np.pi*y))

# "Model": a quadratic function
def sphere(x, a, b):
    return a * (x[0]*x[0] + x[1]*x[1]) + b

def gauss_full(x, sx, sy, amp, b, x0, y0, ang):
    a = np.cos(ang)*np.cos(ang)/2*sx*sx + np.sin(ang)*np.sin(ang)/2*sy*sy
    b = -1*np.sin(2*ang)/4*sx*sx + np.sin(2*ang)/4*sy*sy
    c = np.sin(ang)*np.sin(ang)/2*sx*sx + np.cos(ang)*np.cos(ang)/2*sy*sy
    return amp * np.exp(-1*(a*(x[0]-x0)*(x[0]-x0) + 2*b*(x[0]-x0)*(x[1]-y0) + c*(x[1]-y0)*(x[1]-y0))) + b

def gauss_sb(x, std, b):
    return -1e2 * np.exp(-1*(x[0]*x[0]/2*std*std + x[1]*x[1]/2*std*std)) + b

def gauss_sabx(x, std, amp, b, x0):
    return -1*amp * np.exp(-1*((x[0]-x0[0])*(x[0]-x0[0])/2*std*std + (x[1]-x0[1])*(x[1]-x0[1])/2*std*std)) + b


# Sampled Rastrigin
samp_big_rast = rastrigin(*samp_big)
grid_big_rast = rastrigin(*grid_big)
grid_low_rast = rastrigin(*grid_low)

# Fit the quadratic model to the sample
sopt, scov = curve_fit(sphere, samp_big, samp_big_rast)
grid_big_sphr = sphere(grid_big, *sopt)

# Fit the gaussian model to the sample
gbounds = (0,[1,1e3])
guess = [0.1,1e2]
g0opt, g0cov = curve_fit(gauss_sb, samp_big, samp_big_rast, p0=guess, bounds=gbounds)
print(g0opt)
grid_big_gaus = gauss_sb(grid_big, *g0opt) 

side=(xmax-xmin)/7
samp_big_zoom_x = np.random.random(samp_big_n)*side-side/2
samp_big_zoom_y = np.random.random(samp_big_n)*side-side/2
samp_big_zoom = (samp_big_zoom_x, samp_big_zoom_y)
samp_big_zoom_rast = rastrigin(*samp_big_zoom)

# Fit the gaussian model to the zoomed sample
gbounds = (0,[1e5,1e5])
guess = [1,1e2]
g1opt, g1cov = curve_fit(gauss_sb, samp_big_zoom, samp_big_zoom_rast, p0=guess, bounds=gbounds)
print(g1opt)
grid_big_gaus_zoom = gauss_sb(grid_big, *g1opt)


########################################################################
# Plot set up
########################################################################

# Plots that have their z-values manipulated during the animation
common_props = {"markeredgecolor":"none", "alpha":0.7}
first_sample = None
second_sample = None

def first_init():

    # Grid Rastrigin surface
    ax.plot_surface(*grid_big, grid_big_rast, alpha=0.2, color="magenta");

    # Ground contour
    ax.contour(*grid_low, grid_low_rast, offset=-1, colors="green")

    # Optimum star
    ax.scatter([0], [0], [gap], c="red", s=200, edgecolors="yellow", marker='*')

    # Ground contour around optimum
    ax.contour(*grid_big, grid_big_sphr, offset=-1, colors="red", levels=np.arange(0,1,0.2))

    # Model 3D contour
    #ls = range(0,int(np.ceil(np.max(grid_big_sphr))) +gap,int(np.ceil((np.max(grid_big_sphr) +gap)/cnb)))
    #ax.contour(*grid_big, grid_big_sphr +gap, colors="#ffcc00",levels=ls)

    lg0 = range(0,int(np.ceil(np.max(grid_big_gaus))) +gap,int(np.ceil((np.max(grid_big_gaus) +gap)/cnb)))
    ax.contour(*grid_big, grid_big_gaus +gap, colors="#ffcc00",levels=lg0)

    ax.set_zlim([zmin,zmax])

def second_init():

    # Grid Rastrigin surface
    ax.plot_surface(*grid_big, grid_big_rast, alpha=0.2, color="magenta");

    # Ground contour
    ax.contour(*grid_low, grid_low_rast, offset=-1, colors="green")

    # Optimum star
    ax.scatter([0], [0], [gap], c="red", s=200, edgecolors="yellow", marker='*')

    # Ground contour around optimum
    ax.contour(*grid_big, grid_big_sphr, offset=-1, colors="red", levels=np.arange(0,1,0.2))

    # Model 3D contour
    #ls = range(0,int(np.ceil(np.max(grid_big_sphr))) +gap,int(np.ceil((np.max(grid_big_sphr) +gap)/cnb)))
    #ax.contour(*grid_big, grid_big_sphr +gap, colors="#ffcc00",levels=ls)

    lg1 = range(0,int(np.ceil(np.max(grid_big_gaus_zoom))) +gap,int(np.ceil((np.max(grid_big_gaus_zoom) +gap)/cnb)))
    ax.contour(*grid_big, grid_big_gaus_zoom +gap, colors="#ff5500",levels=lg1)

    ax.set_zlim([zmin,zmax])


def sample_down(points,x,z,i,istart,iend):
    n = i - istart
    nmax = iend - istart
    x,y = x
    for j,p in enumerate(points):
        p.set_data(x[j], y[j])
        z_end = z[j]+gap
        z_dist = ceil - z_end # Distance to ceil
        z_perc = 1 - n/nmax # Remaining percentage of frames
        zj = z_end + z_perc * z_dist
        p.set_3d_properties(zj)

    return points


def first_sample_down(artists, i, istart, iend):
    #ax.scatter(*samp_big, samp_big_rast +gap, c="white",     s=30, edgecolors="none")
    return sample_down(artists, samp_big, samp_big_rast, i, istart, iend)

def second_sample_down(artists, i, istart, iend):
    #ax.scatter(*samp_big_zoom, samp_big_zoom_rast +gap, c="#8888ff", s=30, edgecolors="none")
    return sample_down(artists, samp_big_zoom, samp_big_zoom_rast, i, istart, iend)


# animation function.  This will be called sequentially with the frame number
def first_animation(i):
    global first_sample, second_sample

    print("{}/{}".format(i,nb_frames), flush=True)

    first_sample = sum([ax.plot([], [], [], 'o', c="white", **common_props) for k in range(samp_big_n)], [])
    artists = first_sample_down(first_sample, i, 0, nb_frames//2)

    # Rotate the view point around
    #ax.view_init(view_alt, view_angle + np.degrees(i*(2*np.pi/nb_frames)))
    fig.canvas.draw()

    return artists


# animation function.  This will be called sequentially with the frame number
def second_animation(i):
    global first_sample, second_sample

    print("{}/{}".format(i,nb_frames), flush=True)

    second_sample = sum([ax.plot([], [], [], 'o', c="#8888ff" , **common_props) for k in range(samp_big_n)], [])
    artists = second_sample_down(second_sample, i, nb_frames//2, nb_frames)

    # Rotate the view point around
    #ax.view_init(view_alt, view_angle + np.degrees(i*(2*np.pi/nb_frames)))
    fig.canvas.draw()

    return artists


# instantiate the animator.

# Save as mp4. This requires mplayer or ffmpeg to be installed
#anim.save('illuia.mp4', fps=15, extra_args=['-vcodec', 'libx264'])

#fig,ax = init()
#first_anim = animation.FuncAnimation(fig, first_animation, init_func=first_init, frames=nb_frames, interval=30)#, blit=True)
#first_anim.save('illuia_0.gif',writer='imagemagick',fps=20);
#plt.show()
#
#plt.clf()
#
#fig,ax = init()
#second_anim = animation.FuncAnimation(fig, second_animation, init_func=second_init, frames=nb_frames, interval=30)#, blit=True)
#second_anim.save('illuia_1.gif',writer='imagemagick',fps=20);
#plt.show()


fig,ax = init()
for i in range(nb_frames//2):
    first_init()
    first_animation(i)
    fig.savefig(tmpfile.format(i)+'.png'.format(i))#, transparent = True)
    ax.cla()
    ax.clear()

for i in range(nb_frames//2, nb_frames):
    second_init()
    second_animation(i)
    fig.savefig(tmpfile.format(i)+'.png'.format(i))#, transparent = True)
    ax.cla()
    ax.clear()

args = ["/usr/bin/convert", "-delay", "10", "-loop" , "0", "-dispose", "Background", outfile+"*.png", outfile+".gif"]
print(" ".join(args))
#subprocess.call(args, shell=True)