from tkinter import *
from random import *
from time import *

# constants for health states of individuals
HEALTHY = 0
INFECTED = 1
DAY1 = 2
DAY2 = 3
DAY3 = 4
DAY4 = 5
IMMUNE = 6

def display(c, matrix):
# display current matrix of people on canvas c
    for row in range(len(matrix)):
        for col in range(len(matrix[row])):
            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, outline = "black")

def test_display():
    # create a canvas of size 200 X 200
    window = Tk()
    c = Canvas(window, width=200, height=200)
    c.pack()
    # initialize population matrix (list of lists) randomly
    matrix = []
    for i in range(0,20):
        row = []
        for j in range(0,20):
            row.append(randint(0,6))
        matrix.append(row)
    # display the population matrix using your display function
    display(c, matrix)
    c.update()
    window.mainloop()    # wait until user clicks close button on window
    return None

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 nextday(matrix):
    # create new matrix, initialized to all zeroes
    newmatrix = []
    for i in range(0,20): 
        newmatrix.append([0] * 20)
    # create next day
    for i in range(0,20):
        for j in range(0,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(0,4):  # repeat 4 times
                    # pick a random person at position random x and y
                    random_x = randint(0,19)
                    random_y = randint(0,19)
                    if contagious(matrix, random_x, random_y):
                        newmatrix[i][j] = INFECTED
    return newmatrix

def run_simulation():
    # create a canvas of size 200 X 200
    window = Tk()
    c = Canvas(window, width=200, height=200)
    c.pack()
    # initialize matrix to all healthy individuals
    matrix = []
    for i in range(0,20): 
        matrix.append([0] * 20)
    # infect one random person
    matrix[randint(0,19)][randint(0,19)] = INFECTED
    display(c, matrix)
    c.update()
    # run the simulation for 15 "days"
    sleep(1)
    for day in range(1,16):
        print("Day ", day)
        matrix = nextday(matrix)
        display(c, matrix)
        c.update()
        sleep(1)
    window.mainloop()   # wait until user clicks close on window
    return None