import tkinter
from tkinter import Canvas
from random import randrange
from time import sleep


# Constants for health states of an individual

HEALTHY = 0
INFECTED = 1
DAY1 = 2
DAY2 = 3
DAY3 = 4
DAY4 = 5
IMMUNE  = 6



def display(matrix,c):
    for row in range(len(matrix)):
        for col in range(len(matrix[0])):
            person = matrix[row][col]
            if person == HEALTHY:
                color = "white"  
            elif person == INFECTED:
                color = "pink"
            elif person >= DAY1 and person <= DAY4:
                color = "red"
            else:          # non-contagious or wrong input
                color = "purple"
            c.create_rectangle(col*10, row*10, col*10 + 10, row*10 + 10, fill = color) 
            
def test_display():
    window = tkinter.Tk()
    # create a canvas of size 200 X 200
    c = Canvas(window,width=200,height=200)
    c.pack()
    # create a randomly filled matrix
    matrix  = []
    for i in range(20):
        row = []
        for j in range(20):
            row.append(randrange(7))
        matrix.append(row)
    # display the matrix using your display function
    display(matrix,c)
   


def immune(matrix, i, j):
    return matrix[i][j] == IMMUNE

def contagious(matrix, i, j):
    return matrix[i][j] >= DAY1 and matrix[i][j] <= DAY4

def infected(matrix, i, j):
    return matrix[i][j] == INFECTED

def healthy(matrix, i, j):
    return matrix[i][j] == HEALTHY 



def update(matrix):
    # create new matrix, initialized to all zeroes
    newmatrix = []
    for i in range(20):
        newmatrix.append([0] * 20)
    # create next day
    for i in range(20):
        for j in range(20):
            if immune(matrix, i, j):
                newmatrix[i][j] = IMMUNE 
            elif infected(matrix, i, j) or  contagious(matrix, i, j): 
                newmatrix[i][j] = matrix[i][j] + 1
            elif healthy(matrix, i, j):
                for k in range(4):  # repeat 4 times
                    if contagious(matrix, randrange(20),randrange(20)):
                        newmatrix[i][j] = INFECTED
    return newmatrix

def test_update():
    window = tkinter.Tk()
    # create a canvas of size 200 X 200
    c = Canvas(window,width=200,height=200)
    c.pack()   
    # initialize matrix a to all healthy individuals
    matrix= []
    for i in range(20):
        matrix.append([0] * 20)
    # infect one random person
    matrix[randrange(20)][randrange(20)] = INFECTED
    # display the initial matrix 
    display(matrix,c)
    sleep(0.3)
    # run the simulation for 10 days
    for day in range(0, 10):
        c.delete(tkinter.ALL)
        matrix = update(matrix)
        display(matrix,c)
        sleep(0.3)
        c.update() #force new pixels to display