import cmocean as cmo
import numpy as np
import scipy.constants as spc
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.animation
from ipywidgets import interactive, fixed
from IPython.display import display, Math, HTML
from matplotlib.animation import FuncAnimation

pi = spc.pi
%matplotlib inline
%config InlineBackend.figure_format = 'retina'
plt.rcParams.update(
    {
        'mathtext.fontset':'cm', 
        'font.family':'serif', 
        'font.sans-serif':'Times New Roman', 
        'figure.figsize':[6,6], 
        'figure.titlesize':20,
        'figure.edgecolor': 'w',
        'figure.facecolor': 'w',
        "figure.dpi":90, 
        'savefig.dpi':1200,
        'axes.titlesize':20, 
        'axes.labelsize':18, 
        'axes.titley': 1.0, 
        'axes.titlepad': 5.0,
        'axes.grid': True,
        'axes.grid.which': 'both',
        'axes.linewidth': 2,
        'grid.alpha': 1/5,
        'axes.grid.axis': 'both',
        'grid.linestyle': '--',
        'xtick.labelsize':14,'ytick.labelsize':14,
        'xtick.major.size':6,'ytick.major.size':6,
        'xtick.major.width':1.25, 'ytick.major.width':1.25, 
        'xtick.direction':'inout','ytick.direction':'inout',
        'xtick.top':False, 'ytick.right':False,
        'legend.title_fontsize':14, 'legend.fontsize':14,
        'legend.borderaxespad': 1, 'legend.borderpad': 0.5,
        'legend.framealpha': 1,
        'legend.handleheight': 0.5, 'legend.handlelength': 2.0, 'legend.handletextpad':0.5,
        'legend.labelspacing': 0.25,
        'legend.fancybox':False,
        'legend.edgecolor': '0',
        'legend.frameon': True,
        'legend.markerscale': 1.25,
        'animation.embed_limit':2**128,
        'animation.html': 'jshtml'
    }
)

\begin{align} \mathbf{r}_{i}^{n+1} &= \mathbf{r}_{i}^{n} + \mathbf{v}_{i}^{n} \cdot \Delta t\\\\ \mathbf{v}_{i}^{n} &= v_{0} \begin{pmatrix}\cos(\theta^{n}_{i}) \\ \sin(\theta^{n}_{i})\end{pmatrix}\\\\ \theta^{n+1}_{i} &= \langle \theta^{n} \rangle_{i, R} + \eta^{n}_{i} \end{align}

• $\eta$ adds Gaussian noise that ensures the model doesn't become linearly stable.
• $\langle \theta^{n} \rangle_{i, R}$ is a nearest neighbor average that takes into account all neighbors within a radius of $R$.

def r_hat(theta):
    return np.array([np.cos(theta), np.sin(theta)])

N  = 100
L  = 8.0
R  = 1.0
η  = 0.5
v0 = 1.5
Δt = 0.1
Nt = 100

x, y, vx, vy, Θ = np.zeros((5, Nt+1, N, 1))
t = np.arange(Nt+1)*Δt

np.random.seed(17)
x[0], y[0], Θ[0] = np.random.rand(3, N, 1)
x[0] *= L
y[0] *= L
Θ[0] *= 2*pi
vx[0], vy[0] = v0*r_hat(Θ[0])

for n in range(Nt):
    x[n+1] = (x[n] + vx[n]*Δt)%L
    y[n+1] = (y[n] + vy[n]*Δt)%L
    
    avgϴ = ϴ[n]
    for i in range(N):
        neighbors = (x[n]-x[n,i])**2+(y[n]-y[n,i])**2 < R**2
        sx = np.sum(np.cos(ϴ[n][neighbors]))
        sy = np.sum(np.sin(ϴ[n][neighbors]))
        avgϴ[i] = np.arctan2(sy, sx)
    ϴ[n+1] = avgϴ + η*(np.random.rand(N,1) - 0.5)
    vx[n+1] = v0 * np.cos(ϴ[n+1])
    vy[n+1] = v0 * np.sin(ϴ[n+1])
    
fig, [a,b,c,d] = plt.subplots(1, 4, figsize = (16, 4), tight_layout = True)
a.quiver(x[0],y[0],vx[0],vy[0], color = 'k')
a.set(xlim = (0, L), ylim = (0, L))
a.axis('off')
b.quiver(x[10],y[10],vx[10],vy[10], color = 'b')
b.set(xlim = (0, L), ylim = (0, L))
b.axis('off')
c.quiver(x[20],y[20],vx[20],vy[20], color = 'm')
c.set(xlim = (0, L), ylim = (0, L))
c.axis('off')
d.quiver(x[-1],y[-1],vx[-1],vy[-1], color = 'r')
d.set(xlim = (0, L), ylim = (0, L))
d.axis('off')
plt.show()

# frames = 48
# fps = 24

# fig, ax = plt.subplots(1, 1, figsize = (6, 6))
# field = plt.quiver(x[0], y[0], vx[0], vy[0], color = 'k')
# def animate(i, Q, X, Y, U, V):
#     Q.set_offsets(np.hstack([X[i], Y[i]]))
#     Q.set_UVC(U[i], V[i])
#     return Q,
# ax.set(xlim = (0, L), ylim = (0, L))
# ax.axis('off')
# plt.close()

# anim = FuncAnimation(fig, 
#                      animate, 
#                      fargs = (field, x, y, vx, vy), 
#                      interval = int(1e3/fps), 
#                      frames = np.arange(0, Nt, int(Nt/frames)), 
#                      blit = False)
# anim

# %matplotlib inline
# %config InlineBackend.figure_format = 'retina'
# writergif = mpl.animation.PillowWriter(fps = 24) 
# anim.save('viscek.gif', writer = writergif)