import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider
from scipy.integrate import odeint


def Lissajou(X, t):
    
    [x, vx, y, vy] = X
    return [vx, -sx.val ** 2 * x - sf.val * vx, vy, -sy.val ** 2 * y - sf.val * vy]

def tracer(event):
    
    [X, VX, Y, VY] = np.transpose(odeint(Lissajou, [x0, vx0, y0, vy0], np.linspace(0, sT.val, 1000)))
    trajectoire.set_data(X, Y)
    plt.draw()

def click(event):
    
    global x0, vx0, y0, vy0, selectPos, arrow
    
    if str(event.inaxes)[:11] == 'AxesSubplot':
        if selectPos:
            x0 = event.xdata
            y0 = event.ydata
            pt.set_data(x0, y0)
            selectPos = False
            plt.draw()
        else:
            vx0 = event.xdata - x0
            vy0 = event.ydata - y0
            pt.set_data([], [])
            arrow.remove()
            arrow = ax.arrow(x0, y0, vx0, vy0, width = .02, fill = False)
            selectPos = True    
            tracer(1)
    

x0, y0, vx0, vy0 = 0, 1, 1, 0
    
selectPos = True

fig = plt.figure()
ax = fig.add_subplot(121)
ax.set_xlim(-2, 2)
ax.set_ylim(-2, 2)
ax.grid(True)
ax.set_aspect('equal')

pt, = ax.plot([], [], 'bo', color = 'red')
trajectoire, = ax.plot([], [])

fig.canvas.mpl_connect('button_press_event', click)

axOmegaX = fig.add_axes([.6, .6, .3, .04])
sx = Slider(axOmegaX, r'$\omega_x$', .1, 3, 1.2)
sx.on_changed(tracer)

axOmegaY = fig.add_axes([.6, .55, .3, .04])
sy = Slider(axOmegaY, r'$\omega_y$', .1, 3, 1)
sy.on_changed(tracer)

axf = fig.add_axes([.6, .4, .3, .04])
sf = Slider(axf, r'$f$', 0, .1, 0)
sf.on_changed(tracer)


axT = fig.add_axes([.6, .3, .3, .04])
sT = Slider(axT, r'$T_{max}$', 0.1, 200, 10)
sT.on_changed(tracer)

arrow = ax.arrow(x0, y0, vx0, vy0, width = .02, fill = False)
tracer(1)

mng = plt.get_current_fig_manager()       # Vous pouvez décommenter ça si vous utilisez 'TkAgg'
mng.window.state('zoomed')
plt.show()