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

# creates the four colors for our simulation, EMPTY is white,
# TREE is green, FIRE is red and BURNT is black
colors = cc.ColorConverter.colors
cols = [colors['w'],
        colors['g'],
        colors['r'],
        colors['k']]
ff1 = ListedColormap(cols)

# TODO revise this method to include a boolean parameter called torus
#  when true, the function should wrap around the edges, when false
#  return the function as written.
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

# TODO write this method to include all eight neighbors. Use the 
#  torus parameter as described above
def moore_neighbors(i, j, size, torus):
    pass


class Forest():
    
    states = {"EMPTY":0, "TREE":1, "FIRE":2, "BURNT":3}
    
    def __init__(self, size):
        self.size = size
        self.x = np.arange(size)
        self.y = np.arange(size)
        self.cells = np.zeros((size, size))
        
    def set_fire(self, locs):
        for p in locs:
            self.cells[p[0] + self.size / 2, p[1] + self.size / 2] = Forest.states["FIRE"]
        
    def random_setup(self, d):
        for i in range(self.size):
            for j in range(self.size):
                if (np.random.random() < d):
                    self.cells[i, j] = Forest.states["TREE"]
                else:
                    self.cells[i, j] = Forest.states["EMPTY"]
        
    def image_setup(self):
        plt.title("Forest Fire Simulation")
        self.plt = plt.imshow(self.cells, interpolation='nearest', 
                            origin='bottom', 
                            vmin=np.min(Forest.states["EMPTY"]),
                            vmax=np.max(Forest.states["BURNT"]), 
                            cmap=ff1)

    def update(self):
        newcells = np.zeros((self.size, self.size))
        for i in range(self.size):
            for j in range(self.size):
                neighbors = von_neuman_neighbors(i, j, self.size)
                
                # TODO Add in code to make the fire spread. If a tree is near 
                #  a fire cell, it catches. All trees on fire burn out completely in
                #  one timestep.
                if self.cells[i,j] == Forest.states["FIRE"]:
                    newcells[i,j] = self.cells[i,j]
                elif self.cells[i,j] == Forest.states["TREE"]:
                    newcells[i,j] = self.cells[i,j]
                elif self.cells[i,j] == Forest.states["BURNT"]:
                    newcells[i,j] = self.cells[i,j]
                elif self.cells[i,j] == Forest.states["EMPTY"]:
                    newcells[i,j] = self.cells[i,j]
                    
                # TODO Revise to have burnt trees turn into empty cells randomly
                
                # TODO Revise to have empty cells grow new trees randomly if they
                #  they are adjacent to a living tree
                    
                    
        # TODO Add in a probability of a lightning strike each round. If it hits a 
        #  tree, the tree is set on fire.
        self.cells = newcells
        
    def plot(self):
        self.plt.set_data(self.cells)
        return self.plt
        
size = 60
density = 0.35

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

middle = ((0, 0), )

ff = Forest(size)
ff.random_setup(density)
ff.set_fire(middle)

################
# ANIMATION
ff.image_setup()

def update(data):
    ff.update()
    return ff.plot(),

def data_gen():
    while True: yield 1

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