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

#returns True if person (matrix[r][c]) is healthy
def healthy(matrix, r, c):
    return matrix[r][c] == HEALTHY 

#returns True if person (matrix[r][c]) is infected
def infected(matrix, r, c):
    return matrix[r][c] == INFECTED

#returns True if person (matrix[r][c]) is contagious
def contagious(matrix, r, c):
    return matrix[r][c] >= DAY1 and matrix[r][c] <= DAY4

#returns True if person (matrix[r][c]) is immune
def immune(matrix, r, c):
    return matrix[r][c] == IMMUNE

# takes the current status of the population and returns the next day's data
def simNextDay(data):
   nextDayData = []      # create new matrix and initialize
   for i in range(20): 
      nextDayData.append([0] * 20)    
   for i in range(20):   # create next day
      for j in range(20):
         if immune(data, i, j):
            nextDayData[i][j] = IMMUNE 
         elif infected(data,i,j) or contagious(data,i,j): 
            nextDayData[i][j] = data[i][j] + 1
         elif healthy(data, i, j):
            nextDayData[i][j] = meetPeople(data, i, j)
   return nextDayData


# dependin on randomly chosen people's status returns INFECTED
# or if not infected then it returns the current status of the person back
def meetPeople(currMatrix, row, col): 
   for counter in range(4):  # repeat 4 times
       if contagious(currMatrix, randrange(20), randrange(20)):
          return INFECTED
   return currMatrix[row][col]

# displays the matrix data with appropriate colors
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)

# Just to test the display function
def test_display():
    window = tkinter.Tk()   
    # create a canvas of size 200 X 200
    c = Canvas(window,width=200,height=200)
    c.pack()  
    matrix  = []
    # create a randomly filled matrix
    for i in range(20):
        row = []
        for j in range(20):
            row.append(randrange(7)) # random status of illness
        matrix.append(row)
    # display the matrix using your display function
    display(matrix,c)


# simulates the spread of flu for numOfDays
def simulateFlu(numOfDays):
    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
    population = []
    for i in range(20):
        population.append([0] * 20)

    # infect one random person
    population[randrange(20)][randrange(20)] = INFECTED
    display(population, c)
    sleep(0.3)  # Wait in order to show 1st infection

    # run the simulation for required num of days
    for day in range(0, numOfDays):
        c.delete(tkinter.ALL)
        population = simNextDay(population)
        display(population,c)
        sleep(0.3)   # Wait in order to show change
        c.update()  #Force changes to display -update screen