#! /usr/bin/env python3
#
def forest_fire ( m, n ):

  import numpy as np
#
#  Initialize the height of the trees in the forest.
#
  forest = np.random.randint ( low = 0, high = 5, size = [ m, n ] )
#
#  One square is on full fire.
#
  i = np.random.randint ( low = 0, high = m )
  j = np.random.randint ( low = 0, high = n )
  forest[i,j] = -4
#
#  Cycle
#
  k = 0
  while ( np.any ( forest < 0 ) and k < 10):
    forest_plot ( forest )
    forest = forest_update ( forest )
    k = k + 1
#
#  Terminate.
#
  print ( '' )
  print ( 'forest_fire():' )
  print ( '  Normal end of execution.' )

  return

def forest_plot ( forest ):

  import matplotlib.pyplot as plt

  m, n = forest.shape

  plt.clf ( )
  plt.axis ( 'equal' )
  plt.axis ( 'off' )

  outline = True
  for i in range ( 0, m ):
    for j in range ( 0, n ):

      if ( forest[i,j] == 0 ):
        rgb = [1,1,1]
      elif ( forest[i,j] == 1 ):
        rgb = [0.9,1.0,0.9]
      elif ( forest[i,j] == 2 ):
        rgb = [0.7,1.0,0.7]
      elif ( forest[i,j] == 3 ):
        rgb = [0.5,1.0,0.5]
      elif ( forest[i,j] == 4 ):
        rgb = [0.3,1.0,0.3]
      elif ( forest[i,j] == -1 ):
        rgb = [1.0,0.9,0.9]
      elif ( forest[i,j] == -2 ):
        rgb = [1.0,0.7,0.7]
      elif ( forest[i,j] == -3 ):
        rgb = [1.0,0.5,0.5]
      elif ( forest[i,j] == -4 ):
        rgb = [1.0,0.3,0.3]

      x = [ j, j, j+1, j+1, j ]
      y = [ i, i+1, i+1, i, i ]
      plt.fill ( x, y, color = rgb )
      plt.plot ( x, y, 'k-' )

  plt.title ( 'Forest Fire' )
  filename = 'forest_fire.png'
  plt.savefig ( filename )
  print ( '  Graphics saved as "' + filename + '"' )
  plt.show ( )
  plt.close ( )

  return

def forest_update ( forest ):

  import numpy as np

  m, n = forest.shape

  new_forest = forest.copy ( )
  
  for i in range ( 0, m ):
    for j in range ( 0, n ):

      if ( forest[i,j] < 0 ):

        new_forest[i,j] = forest[i,j] + 1

      elif ( 0 < forest[i,j] ):
        
        im1 = ( i - 1 ) % m
        ip1 = ( i + 1 ) % m
        jm1 = ( j - 1 ) % n
        jp1 = ( j + 1 ) % n

        s = ( forest[im1,j] < 0 ) + ( forest[ip1,j] < 0 ) \
          + ( forest[i,jm1] < 0 ) + ( forest[i,jp1] < 0 )

        ignite = np.random.random ( )

        if ( ignite < 0.50 * s ):
          new_forest[i,j] = - new_forest[i,j]
 
  forest = new_forest.copy ( )

  return forest

if ( __name__ == "__main__" ):
  forest_fire ( 10, 10 )