import copy , time , pygame , sys , random

# Get things up and running

pygame.init()

# Create a window where stuff can happen

my_screen = pygame.display.set_mode( ( 800 , 600 ) , 5 )

# Load a font so that we can write something

#my_font = pygame.font.SysFont( 'monospace' , 36 )

# Our grid's width and height

grid_width = 100
grid_height = 50

# The position and size of the boxes used to display the grid

box_width = 8
box_height = 12

grid_x = 0
grid_y = 0

# Colour of the boxes on the grid

r_color = (  10 ,  10 ,  10 ) # Colour for Removed
s_color = (   0 ,   0 , 200 ) # Colour for Susceptible
i_color = ( 200 ,   0 , 100 ) # Colour for Infectious
w_color = ( 100 , 100 , 100 ) # Colour for Wall

# Beta, gamma and mu are the contact rate, the infectiousness and recovery rate of our disease

# A cold due to renewed susceptibility

beta = 0.4
gamma = 0.6
mu = 0.6

# Initialize grid

grid = []

# This bit makes sure that our grid is the right size

if len( grid ) > grid_height:
    grid_height = len( grid )
for row in grid:
    if len( row ) > grid_width:
        grid_width = len( row )

if len( grid ) < grid_height:
    grid+= [ [] for i in range( grid_height - len( grid ) ) ]

for row in grid:
    if len( row ) < grid_width:
        row+= [ 1 for i in range( grid_width - len( row ) ) ]

# Create a firewall

firewall_x = int( grid_width / 3 )

for row in grid:
    row[ firewall_x ] = 3

# Now leave some holes

porosity = 0.2
num_holes = int( grid_height * porosity )
hole_count = 0

while hole_count < num_holes:
    hole_y = random.randint( 0 , grid_height - 1 )
    if grid[ hole_y ][ firewall_x ] == 3:
        grid[ hole_y ][ firewall_x ] = 1
        hole_count += 1

# Infect a few random people

to_infect = 10

for i in range( to_infect ):
    cell_available = False
    while not( cell_available ):
        x = random.randint( int( grid_width / 2 ) , grid_width - 1 )
        y = random.randint( int( grid_height / 2 ) , grid_height - 1 )
        cell_available = ( grid[ y ][ x ] == 1 )
    grid[ y ][ x ] = 2

# How to find the neighbours of a cell...

neighbor_list = []

for j in [ -1 , 0 , 1 ]:
    for i in [ -1 , 0 , 1 ]:
        if not( i == 0 and j == 0 ):
            neighbor_list.append( ( i , j ) )

# This creates a loop that only ends when you close the window

still_looping = True

while still_looping:

    # This looks to see whether the close button has been pressed

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            still_looping = False

    # Clean the board before we start!!

    my_screen.fill( ( 100 , 100 , 150 ) )

    # Reset counters for the state of the system

    S = 0
    I = 0
    R = 0

    # Create a working copy of our grid so that we can iterate

    ref_grid = copy.deepcopy( grid )

    # Now let's iterate

    for y in range( grid_height ):
        for x in range( grid_width ):

            # Start by displaying the current state

            cell_val = ref_grid[ y ][ x ]

            if cell_val == 0:
                my_color = r_color
            if cell_val == 1:
                my_color = s_color
            if cell_val == 2:
                my_color = i_color
            if cell_val == 3:
                my_color = w_color

            pygame.draw.rect( my_screen , my_color , ( grid_x + box_width * x , grid_y + box_height * y , box_width , box_height ) , 0 )

            # Counter for neighbours

            neighbor_count = 0

            # Work through neighbours

            for ( i , j ) in neighbor_list:
                neighbor_x = x + i
                neighbor_y = y + j

                # As long as they are on the grid!

                if neighbor_x >= 0 and neighbor_y >= 0 and neighbor_x < grid_width and neighbor_y < grid_height:

                    # If so then count whether they are infected
                    if ref_grid[ neighbor_y ][ neighbor_x ] == 2:
                        neighbor_count = neighbor_count + 1

            # Then iterate according to our rules

            # If dead/removed then we should roll a die to come to life

            if cell_val == 0:
                R += 1
                if random.random() < mu:
                    cell_val = 1

            # If alive with an infected neighbour then roll a die to infect

            if cell_val == 1 and neighbor_count >0:
                S += 1
                if random.random() < beta:
                    cell_val = 2

            # If infected then roll a die to die

            if cell_val == 2:
                I += 1
                if random.random() < gamma:
                    cell_val = 0

            # Update the grid

            grid[ y ][  x ] = cell_val

    # Show off our handiwork

    pygame.display.update()

    # Remind us of the values of the parameters and the current state

    print( '\n\n beta = ' , beta , '\t\t S = ' , S)
    print( 'gamma = ' , gamma , '\t\t I = ' , I )
    print( '   mu = ' , mu , '\t\tR = ', R , '\n\n' )

# Finish things off
pygame.quit()
sys.exit()