import numpy as np
import random
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.colors import LogNorm
import matplotlib.colors as cc
from matplotlib.colors import ListedColormap

WOLFCOUNT = 10
SHEEPCOUNT = 30

WOLFENERGYSTART = 200.0
WOLFENERGY = 10
WOLFMOVECOST = 1
WOLFLINGLEVEL = 50

SHEEPLINGPROB = 0.04

colors = cc.ColorConverter.colors
cols = [colors['w'],
        colors['r']]
acolors = ListedColormap(cols)

def von_neuman_neighbors(i, j, size):
    n = []
    if i > 0:
        n.append((i - 1, j))
    if j > 0:
        n.append((i, j - 1))
    if i < size - 1:
        n.append((i + 1, j))
    if j < size - 1:
        n.append((i, j + 1))
    return n

def moore_neighbors(i, j, size):
    n = []
    if i > 0:
        n.append((i - 1, j))
        if j > 0:
            n.append((i - 1, j - 1))
        if j < size - 1:
            n.append((i - 1, j + 1))
    if j > 0:
        n.append((i, j - 1))
    if i < size - 1:
        n.append((i + 1, j))
        if j > 0:
            n.append((i + 1, j - 1))
        if j < size - 1:
            n.append((i + 1, j + 1))
    if j < size - 1:
        n.append((i, j + 1))
    return n

class Agent():
    
    def __init__(self, x, y, world, color):
        self.x = x
        self.y = y
        self.world = world
        self.color = color
        
    def wander(self):
        n = von_neuman_neighbors(self.x, self.y, self.world.size)
        self.x, self.y = random.choice(n)
        
    def update(self):
        pass

class Wolf(Agent):
    
    def __init__(self, x, y, world, energy):
        Agent.__init__(self, x, y, world, 1)
        self.energy = energy

    def update(self):
        self.wander()
        self.energy -= 0.5
        if self.energy >= WOLFLINGLEVEL:
            self.world.wolves.append(Wolf(self.x, self.y, self.world, self.energy / 2))
            self.energy /= 2
        elif self.energy <= 0:
            self.world.wolves.remove(self)
        else:
            for s in self.world.sheep:
                if self.x == s.x and self.y == s.y:
                    self.energy += WOLFENERGY
                    self.world.sheep.remove(s)
                    break
        
class Sheep(Agent):
    
    def __init__(self, x, y, world):
        Agent.__init__(self, x, y, world, 0)
    
    def update(self):
        self.wander()
        if (np.random.random() < SHEEPLINGPROB):
            self.world.sheep.append(Sheep(self.x, self.y, self.world))


class CA():
    
    states = {"EMPTY":0, "GRASS":1}
    
    def __init__(self, size):
        self.size = size
        self.x = np.arange(size)
        self.y = np.arange(size)
        self.cells = np.zeros((size, size))
        self.time = 0
        
        self.wolves = []
        for w in range(WOLFCOUNT):
            i, j = np.random.random_integers(0, size, 2)
            self.wolves.append(Wolf(i, j, self, WOLFENERGYSTART))

        self.sheep = []
        for s in range(SHEEPCOUNT):
            i, j = np.random.random_integers(0, size, 2)
            self.sheep.append(Sheep(i, j, self))
        
    def random_setup(self, n):
        for i in range(n):
            a, b = np.random.random_integers(0, size - 1, 2)
            self.cells[a, b] = CA.states["GRASS"]
        
    def image_setup(self):
        
        self.plt = plt.imshow(self.cells, interpolation='nearest', 
                            origin='bottom', 
                            vmin=np.min(0),
                            vmax=np.max(1),
                            cmap=plt.cm.winter)
        ax = []
        ay = []
        ac = []
        for a in self.wolves:
            ax.append(a.x)
            ay.append(a.y)
            ac.append(a.color)
        for f in self.sheep:
            ax.append(f.x)
            ay.append(f.y)
            ac.append(f.color)
        self.ant_plot = plt.scatter(ax, ay, c=ac, s=60, cmap=acolors)

    def update(self):
        for a in self.wolves:
            a.update()
        for a in self.sheep:
            a.update()
        self.time += 1
        
    def plot(self):
        self.plt.set_data(np.transpose(self.cells))
        ap = []
        ac = []
        for a in self.wolves:
            ap.append((a.x, a.y))
            ac.append(a.color)
        for f in self.sheep:
            ap.append((f.x, f.y))
            ac.append(f.color)
        self.ant_plot.set_offsets(ap)
        ac = np.array(ac)
        self.ant_plot.set_array(ac)
        return self.plt, self.ant_plot
        
size = 40

fig, ax = plt.subplots()
ax.set_ylim(-1, size)
ax.set_xlim(-1, size)

ca = CA(size)
ca.image_setup()
#ca.random_setup(100)

numsheep = []
numwolves = []
time = []
       
def update(data):
    ca.update()
    #ca.random_setup(50)
    time.append(ca.time)
    numsheep.append(len(ca.sheep))
    numwolves.append(len(ca.wolves))
    return ca.plot(), 

def data_gen():
    while True: yield 1

ani = animation.FuncAnimation(fig, update, data_gen, blit=False, interval=40)
plt.show()

plt.plot(time, numwolves)
plt.plot(time, numsheep)