#!/usr/bin/env python
# coding: utf-8

# # Spatial evolutionary game simulator
# 
# S. Kolotev, A. Malyutin, E. Burovski, S. Krashakov and L. Shchur, *Dynamic fractals in spatial evolutionary games*, Physica A **499**, 142 (2018)

# In[1]:


import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib import animation

get_ipython().run_line_magic('matplotlib', 'notebook')


# In[2]:


import sys
import cython

print("python ", sys.version)
print("numpy", np.__version__, ", matplotlib", mpl.__version__, ", and Cython", cython.__version__)


# ## Main update function:

# In[3]:


def evolve2(field, b, num_steps=1):
    L = field.shape[0]
    current = np.zeros((L, L), dtype=int)
    scores = np.zeros((L, L), dtype=float)
    
    for step in range(num_steps):
        current = field.copy()
        scores = np.zeros((L, L), dtype=float)
        
        for x in range(L):
            for y in range(L):
                for i in range(-1, 2):
                    for j in range(-1, 2):
                        ix = (x + i) % L
                        jy = (y + j) % L
                        scores[x, y] += (1 - field[ix, jy])
                        
                if field[x, y] == 1:
                    scores[x, y] *= b
        
        for x in range(L):
            for y in range(L):
                bestX = x
                bestY = y
                for i in range(-1, 2):
                    for j in range(-1, 2):
                        ix = (x + i) % L
                        jy = (y + j) % L
                        if (scores[bestX, bestY] < scores[ix, jy]):
                            bestX = ix
                            bestY = jy
                
                field[x, y] = current[bestX, bestY]
    return field


# ### Run the simulation and animate

# In[4]:


L = 75
field = np.zeros((L, L), dtype=int)
field[L//2, L//2] = 1

# draw the initial field
fig = plt.figure()
im = plt.imshow(field, animated=True, cmap=plt.get_cmap('Paired'))

# updater function
def updatefig(*args):
    global field
    field = evolve2(field, 1.81, 1)
    im.set_array(field)
    return im,

# animate!
anim = animation.FuncAnimation(fig, updatefig, frames=200, interval=50, blit=True)

plt.show()


# ### Timings

# In[5]:


L = 44
field = np.zeros((L, L), dtype=int)
field[L//2, L//2] = 1

get_ipython().run_line_magic('timeit', 'evolve2(field, 1.81, 10)')


# # Enter Cython

# In class, we started with the original code and did the conversion iteratively, removing yellow lines from the annotated output. Below is only the final result.

# In[6]:


get_ipython().run_line_magic('load_ext', 'cython')


# In[7]:


get_ipython().run_cell_magic('cython', '-a', '\nimport numpy as np\n\nimport cython\n\n@cython.cdivision(True)\n@cython.boundscheck(False)\ndef evolve2_1(long[:, ::1] field, double b, int num_steps=1):\n    \n    cdef int x, y, L, i, j, ix, jy, step\n    \n    L = field.shape[0]\n    cdef double[:, ::1] scores = np.zeros((L, L), dtype=float)\n    \n    cdef double[:, ::1] _zeros = np.zeros((L, L), dtype=float)\n    cdef long[:, ::1] current = field.copy()\n    \n    for step in range(num_steps):\n        current = field.copy()\n        scores[...] = _zeros\n        \n        for x in range(L):\n            for y in range(L):\n                for i in range(-1, 2):\n                    for j in range(-1, 2):\n                        ix = (x + i) % L\n                        jy = (y + j) % L\n                        scores[x, y] += (1 - field[ix, jy])\n                        \n                if field[x, y] == 1:\n                    scores[x, y] *= b\n        \n        for x in range(L):\n            for y in range(L):\n                bestX = x\n                bestY = y\n                for i in range(-1, 2):\n                    for j in range(-1, 2):\n                        ix = (x + i) % L\n                        jy = (y + j) % L\n                        if (scores[bestX, bestY] < scores[ix, jy]):\n                            bestX = ix\n                            bestY = jy\n                \n                field[x, y] = current[bestX, bestY]\n    return field\n')


# In[10]:


L = 44
field = np.zeros((L, L), dtype=int)
field[L//2, L//2] = 1

get_ipython().run_line_magic('timeit', 'evolve2_1(field, 1.81, 10)')


# In[11]:


345 / 1.98


# The recommended code structure is to have a python-facing `def` function which takes care of the input and output, allocated memory etc, and delegates all heavy lifting to an internal `cdef` function.

# In[ ]:


get_ipython().run_cell_magic('cython', '-a', '\nimport numpy as np\nimport cython\n\ndef evolve2_2(field, b, num_steps=1):\n    # validate input here\n    field = np.atleast_2d(field)\n    # XXX check that field is integer dtype etc\n    \n    # allocate memory here\n    L = field.shape[0]\n    scores = np.zeros((L, L), dtype=float)\n    current = field.copy()\n    _zeros = np.zeros((L, L), dtype=float)\n\n    cdef:\n        int _num_steps = num_steps\n        double _b = b\n    \n    # do the work (`field` array is updated in-place)\n    _evolve2_2_impl(field, _b, _num_steps, scores, _zeros, current)\n    \n    # convert memoryviews to numpy arrays\n    return np.asarray(field)\n\n\n@cython.cdivision(True)\n@cython.boundscheck(False)\ncdef void _evolve2_2_impl(long[:, ::1] field, double b, int num_steps,\n                                  double[:, ::1] scores, double[:, ::1] _zeros,\n                                  long[:, ::1] current) nogil:\n    \n    cdef int x, y, L, i, j, ix, jy, step\n    \n    L = field.shape[0]\n\n    for step in range(num_steps):\n        current[...] = field\n        scores[...] = _zeros\n        \n        for x in range(L):\n            for y in range(L):\n                for i in range(-1, 2):\n                    for j in range(-1, 2):\n                        ix = (x + i) % L\n                        jy = (y + j) % L\n                        scores[x, y] += (1 - field[ix, jy])\n                        \n                if field[x, y] == 1:\n                    scores[x, y] *= b\n        \n        for x in range(L):\n            for y in range(L):\n                bestX = x\n                bestY = y\n                for i in range(-1, 2):\n                    for j in range(-1, 2):\n                        ix = (x + i) % L\n                        jy = (y + j) % L\n                        if (scores[bestX, bestY] < scores[ix, jy]):\n                            bestX = ix\n                            bestY = jy\n                \n                field[x, y] = current[bestX, bestY]\n')


# In[ ]:


L = 44
field = np.zeros((L, L), dtype=int)
field[L//2, L//2] = 1

get_ipython().run_line_magic('timeit', 'evolve2_2(field, 1.81, 10)')


# # Now, for something completely different: `numba` just-in-time compilation

# In[1]:


import numba


# In[21]:


def evolve3(field, b, num_steps=1):
    L = field.shape[0]
    _zeros = np.zeros((L, L), dtype=float)
    _int_zeros = np.zeros((L, L), dtype=int)
    current = _int_zeros.copy()
    scores = _zeros.copy()
    evolve3_impl(field, b, num_steps, _zeros, _int_zeros)
    return field


# Note a single decorator: @jit


@numba.jit(nopython=True)
def evolve3_impl(field, b, num_steps, _zeros, _int_zeros):
    L = field.shape[0]
    current = _int_zeros.copy()
    
    for step in range(num_steps):
        current = field.copy()
        scores = _zeros.copy()
        
        for x in range(L):
            for y in range(L):
                for i in range(-1, 2):
                    for j in range(-1, 2):
                        ix = (x + i) % L
                        jy = (y + j) % L
                        scores[x, y] += (1 - field[ix, jy])
                        
                if field[x, y] == 1:
                    scores[x, y] *= b
        
        for x in range(L):
            for y in range(L):
                bestX = x
                bestY = y
                for i in range(-1, 2):
                    for j in range(-1, 2):
                        ix = (x + i) % L
                        jy = (y + j) % L
                        if (scores[bestX, bestY] < scores[ix, jy]):
                            bestX = ix
                            bestY = jy
                
                field[x, y] = current[bestX, bestY]
    return field


# In[23]:


L = 44
field = np.zeros((L, L), dtype=int)
field[L//2, L//2] = 1

get_ipython().run_line_magic('timeit', 'evolve3(field, 1.81, 10)')


# In[ ]: