from PythonLabs.Canvas import Canvas
from random import *

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

def display(matrix):
    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:
                color = "purple"
            Canvas.Rectangle(col*10,row*10,(col+1)*10,(row+1)*10,fill=color)
    return None

def test_display():
    # create a canvas of size 200 X 200
    Canvas.init(240, 240, "Testing_Display")
    # initialize matrix a randomly
    a = []
    for i in range(0,20):
        row = []
        for j in range(0,20):
            row.append(randint(0,6))
        a.append(row)
    # display the matrix using your display function
    display(a)

def is_immune(matrix, i, j):    # row is i, col is j
    return matrix[i][j] == IMMUNE

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

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

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

def update(matrix):
    # create new matrix for next day, don't change current matrix in place
    newmatrix = []
    for i in range(0,20):
        newmatrix.append([0] * 20)
    # simulate next day
    for i in range(0,20):
        for j in range(0,20):
            if is_immune(matrix,i,j):
                newmatrix[i][j] = IMMUNE
            elif is_infected(matrix,i,j) or is_contagious(matrix,i,j):
                newmatrix[i][j] = matrix[i][j] + 1
            elif is_healthy(matrix,i,j):
                for k in range(0,4):
                    # pick 4 random people
                    random_x = randint(0,19)
                    random_y = randint(0,19)
                    if is_contagious(matrix,random_x,random_y):
                        newmatrix[i][j] = INFECTED
                # if still healthy, newmatrix[i][j] is 0 so this is ok
    return newmatrix

def test_update():
    # create a canvas of size 200 X 200
    Canvas.init(240, 240, "Testing_Update")
    # initialize matrix a to all healthy individuals
    a = []
    for i in range(0,20): 
        a.append([0] * 20)
    # infect one random person
    a[randint(0,19)][randint(0,19)] = INFECTED
    display(a)
    # run the simulation for 15 "days”
    Canvas.delay = 2
    for day in range(0,15):
        a = update(a)
        display(a)
        Canvas.update() # copy pixels to screen and pause