text: # Building evacuation Q-learning algorithm.
  # By Akash Kumar
  # Tutorial: http://mnemstudio.org/path-finding-q-learning-tutorial.htm
  
  import random
  import sys
  
  alpha = 0.1 # learning rate
  gamma = 0.8 # discount factor
  
  # Building graph represented using an adjacency list, where the nodes
  # are the rooms and the edges are the connections between rooms.
  # The last room represents the outside world.
  building = [
      [4],     # room 0
      [3,5],   # room 1
      [3],     # room 2
      [1,2,4], # room 3
      [0,3,5], # room 4
      [1,4,5]  # room 5 (outside)
  ]
  
  # Reward table.
  reward = [
      # action
      [None,None,None,None,-10,None],  # state 0
      [None,None,None,-10,None,100],   # state 1
      [None,None,None,-10,None,None],  # state 2
      [None,-10,-10,None,-10,None],    # state 3
      [-10,None,None,-10,None,100],    # state 4
      [None,-10,None,None,-10,100]     # state 5
  ]
  
  # Q table.
  # Before learning, the agent knows nothing about the environment.
  # So, all Q values are initialized to 0.
  qtable = [
      [0,0,0,0,0,0],
      [0,0,0,0,0,0],
      [0,0,0,0,0,0],
      [0,0,0,0,0,0],
      [0,0,0,0,0,0],
      [0,0,0,0,0,0]
  ]
  
  # A reward function R(s,a) that gives the agent a reward for taking
  # action a in state s.
  def r(s,a):
      return reward[s][a]
  
  # Q function.
  # Used to calculate Q (quality) values for state-action pairs using
  # the Bellman equation and returns the new Q value.
  def q(s,a):
      # The Q value of this state is based on the max Q value of the next state.
      next_state_max_q = max([qtable[a][adj_room] for adj_room in building[a]])
      qtable[s][a] = (qtable[s][a]
          + alpha * (r(s,a) + gamma * next_state_max_q - qtable[s][a]))
  
      return qtable[s][a]
  
  # Used to set all Q values in the Q table to 0.
  def reset_q_table():
      for i in range(len(qtable)):
          for j in range(len(qtable[0])):
              qtable[i][j] = 0
  
  # Used to represent the Q table as a readable string.
  def get_qtable_str():
      output = "[\n"
      for row in qtable:
          output += "\t" + str([round(x,2) for x in row]) + ",\n"
      output += "]\n"
  
      return output
  
  # Returns the index of the maximum element in a number list.
  def get_max_index(l):
      max_i = -1
      max_num = -sys.maxsize
      for i, num in enumerate(l):
          if num > max_num:
              max_num = num
              max_i = i
  
      return max_i
  
  # Used to make the agent learn about its environment for a given number of
  # simulations and learn the values of the Q table.
  def qlearn(num_simulations):
      total_move_count = 0
      for _ in range(num_simulations):
          start_state = random.randrange(len(building))
          end_state = 5
          while start_state != end_state:
              next_state = random.choice(building[start_state])
              q(start_state,next_state)
              start_state = next_state
              total_move_count += 1
  
      avg_move_count = total_move_count / num_simulations
  
      return (f"Completed {num_simulations} Q-learning simulations with an "
          + f"average move count of {avg_move_count}")
  
  # Used to make the agent evacuate the building as fast as possible,
  # after having learnt about its environment.
  def evacuate(num_simulations):
      total_move_count = 0
      for _ in range(num_simulations):
          start_state = random.randrange(len(building))
          end_state = 5
          while start_state != end_state:
              start_state = get_max_index(qtable[start_state])
              total_move_count += 1
  
      avg_move_count = total_move_count / num_simulations
  
      return (f"Completed {num_simulations} evacuations with an average "
          + f"move count of {avg_move_count}")
  
  if __name__ == "__main__":
      num_simulations = 100
      print(qlearn(num_simulations))
      print(evacuate(num_simulations))
      print("Q table:")
      print(get_qtable_str())
      reset_q_table()
  
      num_simulations = 500
      print(qlearn(num_simulations))
      print(evacuate(num_simulations) + "\n")
      reset_q_table()
  
      num_simulations = 1000
      print(qlearn(num_simulations))
      print(evacuate(num_simulations) + "\n")
      reset_q_table()
  