import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import odeint
from matplotlib.animation import FuncAnimation
from functools import partial

#system
def van_der_pol(coords, t, mu, evens, odds, ones):
    x, y = coords*evens, coords*odds
    x_shift = np.roll(x, 1) #for use in computation
    y_prime = mu*(1-x_shift**2)*y-x_shift
    return np.roll(y, -1)+y_prime

#initialize window for blitting
def init(fig, axes, scatter):
    plt.xlim([-3.0, 3.0])
    plt.ylim([-3.0, 3.0])
    return scatter,

#animates one frame
def animate(t, sols, scatter, odds, evens):
    print(t)

    scatter.set_offsets(np.array([np.trim_zeros(sols[t]*evens), np.trim_zeros(sols[t]*odds)]))
    return scatter,

#constants
mu = 0.5

#initial data
M = 10
axis = np.linspace(-3.0, 3.0, M)
init_conds = np.array([(x, y) for x in axis for y in axis])
init_conds = init_conds.reshape(2*M**2)
evens, odds = np.array([1 if i % 2 == 0 else 0 for i in range(2*M**2)]), np.array([1 if i % 2 == 1 else 0 for i in range(2*M**2)]) #will be used to extract the x and y coordinates
ones = np.ones(2*M**2)

#solving
total = 10.0
delta = 0.1
N = int(total/delta)
t = np.linspace(0.0, total, N)
sols = odeint(van_der_pol, init_conds, t, args=(mu, evens, odds, ones))

#setup
fig, axes = plt.figure(), plt.axes(frameon=True)
scatter = plt.scatter([], [], animated=True)

#animation
anim = FuncAnimation(fig, func=animate, frames=N, init_func=partial(init, fig, axes, scatter), blit=True, fargs=(sols, scatter, odds, evens), repeat=False)
anim.save("out.mp4", fps=40)
plt.show()