contagion.py

import random


def main():
    pop = makePopulation(10)
    pop[3][3] = 5
    for i in range(100):
        pop = stepForward(pop)

        # COMMENT AND UNCOMMENT NEXT TWO LINES
        printGrid(pop)
        printBar(pop)


def stepForward(pop):
    newPop = [[-1 for y in range(len(pop))] for x in range(len(pop))]

    for x in range(1, len(pop) - 1):
        for y in range(1, len(pop) - 1):
            if pop[x][y] == -1:
                nCount = sum([
                    n > 0 for n in [
                        pop[x + 1][y], pop[x - 1][y], pop[x][y + 1], pop[x][y -
                                                                            1]
                    ]
                ])
                newPop[x][y] = 5 if random.random() < (0.13 * nCount) else -1

            elif pop[x][y] > 0:
                newPop[x][y] = pop[x][y] - 1

            elif pop[x][y] == 0:
                newPop[x][y] = 0

    return newPop


def safeAcc(grid, x, y):
    try:
        return grid[x][y]
    except:
        return 0


def printGrid(pop):
    for y in pop:
        for x in y:
            if x == -1:
                print('\033[45m ', end='')
            elif x > 0:
                print('\033[44m ', end='')
            elif x == 0:
                print('\033[43m ', end='')
        print("\033[m")
    print()


def printBar(pop):
    infectedCount = sum([sum([1 for y in x if y > 0]) for x in pop])
    recoveredCount = sum([sum([1 for y in x if y == 0]) for x in pop])
    print('-' * recoveredCount, end='')
    print('*' * infectedCount)


def makePopulation(n):
    return [[-1 for y in range(n + 2)] for x in range(n + 2)]


if __name__ == "__main__":
    main()