T = 300 # Kelvin
kb = 1.38e-34
beta = 1/(kb/T)
steps = 2000
size = 150

lattice = np.random.choice([1,-1],size=(size, size))

fig = plt.figure(figsize=(10,10))
im = plt.imshow(lattice, interpolation='none')
title = plt.text(0,0,"Iteration: ")

def init():
    im.set_data(lattice)
    return im, title

def get_spin(m,n):
    global lattice
    if m == 0 or m == lattice.shape[0] or n == 0 or n == lattice.shape[0]:
        return 1
    return lattice[m,n]

def neighbor_sum(m,n, N=size):
    global lattice
    a, b = m, n
    return lattice[(a+1)%N,b] + lattice[a,(b+1)%N] + lattice[(a-1)%N,b] + lattice[a,(b-1)%N]

def Emn( n, m, size=size):
    global lattice
    cum = 0
    s_mn = lattice[n][m]
    for i in range(n-1, n+2):
        for j in range(m-1, m+2):
            if i==n and j == m: continue;
            cum +=  lattice[i %size, j % size]
    return 2*cum*s_mn

def p(emn, beta=beta):
    return np.exp(-1*beta*emn)

def flip_energy(m,n):
    global lattice
    e_mn = Emn(m,n)
    if e_mn <= 0:
        lattice[m][n] *= -1
    elif p(e_mn) > np.random.rand():
        lattice[m][n] *= -1

def animate(i):
    global lattice
    m,n = np.random.randint(0,lattice.shape[0]-1, size=(2,))
    flip_energy(m,n)
    im.set_data(lattice)

def batch_update(i):
    for t in range(200):
        animate(i)
    title.set_text(f"Iteration: {i}")
    if (i%100 == 0):
        print(f"Iteration: {i}")
    return im, title

anim = animation.FuncAnimation(fig, batch_update, init_func=init,
                               frames=steps, interval=100)
print("Saving animation...")
anim.save('ising model.mp4',fps=30, extra_args=['-vcodec', 'libx264'])
print("Done :D")
# HTML(anim.to_html5_video())