text: # coding=utf-8
  from random import random, randint
  
  import click as cl
  import numpy as np
  
  state_grid = np.array([[0, 0, 0, 0, 0, 0, 0, 0, 1],
                         [0, 0, 0, 0, 0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0, 0, 0, 0, 0],
                         [-1, -1, -1, -1, -1, -1, -1, -1, 0],
                         [0, 0, 0, 0, 0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0, 0, 0, 0, 0]])
  
  updated_grid = np.array([[0, 0, 0, 0, 0, 0, 0, 0, 1],
                           [0, 0, 0, 0, 0, 0, 0, 0, 0],
                           [0, 0, 0, 0, 0, 0, 0, 0, 0],
                           [0, -1, -1, -1, -1, -1, -1, -1, -1],
                           [0, 0, 0, 0, 0, 0, 0, 0, 0],
                           [0, 0, 0, 0, 0, 0, 0, 0, 0]])
  
  
  def make_transition_matrix(matrix):
      """Derive a transition matrix from a 2D array representing a grid world.
  
       By convention, the four transitions (up, right, down, and left) are represented by columns 0 to 3, respectively.
       The value -1 represents an impassible value. The value 1 represents the reward value.
  
      Parameters
      ----------
      matrix : ndarray of ints
          A 2D array representing all possible transitions.
  
      Returns
      -------
      transition_matrix : ndarray of ints
          An N X 4 2D array where N is the total size of matrix.
  
      Notes
      -----
      It would be possible to remove the impassible values from the transition matrix but I chose not to for ease of
      indexing.
      """
      padded_matrix = np.pad(matrix, pad_width=1, mode="constant", constant_values=-1)
  
      up = np.roll(padded_matrix, shift=1, axis=0)[1:-1, 1:-1].flatten()
      right = np.roll(padded_matrix, shift=-1, axis=1)[1:-1, 1:-1].flatten()
      down = np.roll(padded_matrix, shift=-1, axis=0)[1:-1, 1:-1].flatten()
      left = np.roll(padded_matrix, shift=1, axis=1)[1:-1, 1:-1].flatten()
  
      transition_matrix = np.stack((up, right, down, left), axis=0).T
  
      return transition_matrix
  
  
  def index_1d(row, col):
      """For a given row and column in a grid world, return the row that cell would appear on in the transition matrix.
  
      Parameters
      ----------
      row : int
          The row number. Indexing starts at 0.
      col : int
          The column number. Indexing starts at 0.
  
      Returns
      -------
      transition_row : int
          The row on the transition matrix the row and column would appear on.
      """
      max_cols = 9
      return row * max_cols + col
  
  
  def reverse_index_1d(state):
      """The inverse of index_1d. From an row number in the transition matrix return the row and column in the grid world.
  
      Parameters
      ----------
      state : int
          An integer representing the row in the transition matrix.
  
      Returns
      -------
      row : int
          The row in the grid world.
      column : int
          The column in the grid world.\
      """
      max_cols = 9
      row = state / max_cols
      col = state % max_cols
      return row, col
  
  
  def peek_next_state(grid, state, action):
      """For a given grid, state, and action return the next state in the transition matrix.
  
      If the next state is less than 0, stay in the same state.
  
      Parameters
      ----------
      grid : ndarray of ints
          The transition matrix.
      state : int
          The current state of the transition matrix (represented as a row number)
      action : int
          The action being observed. 0 = up. 1 = right. 2 = down. 3 = right.
  
      Returns
      -------
      next_state : int
          The next state in the transition matrix.
      """
      assert state <= grid.shape[0], "State larger than grid size"
      assert action <= grid.shape[1], "Action larger than grid size"
  
      if grid[state, action] < 0:
          return state
      else:
          row, col = reverse_index_1d(state)
          if action == 0:
              return index_1d(row - 1, col)
          elif action == 1:
              return index_1d(row, col + 1)
          elif action == 2:
              return index_1d(row + 1, col)
          elif action == 3:
              return index_1d(row, col - 1)
          else:
              raise ValueError("Action cannot be greater than 3")
  
  
  def peek_reward(grid, state, action):
      """For a given grid, state, and action return the reward in the transition matrix.
  
      Parameters
      ----------
      grid : ndarray of ints
          The transition matrix.
      state : int
          The current state of the transition matrix (represented as a row number)
      action : int
          The action being observed. 0 = up. 1 = right. 2 = down. 3 = right.
  
      Returns
      -------
      reward : int
          The reward in the transition matrix
      """
      return grid[state, action]
  
  
  def update_q(q, state, action, alpha, gamma, reward, next_state):
      """Update the Q matrix according to formula:
  
      .. math::Q(s, a) ← Q(s, a) + α r + γ max a Q(s 0 , a) − Q(s, a)
  
  
      Parameters
      ----------
      q : ndarray of floats
          The 2D matrix the represents the Q function.
      state : int
          The current state of the transition matrix.
      action : int
          The action taken.
      alpha : float
          The learning rate. This determines how much each step is weighted.
      gamma : float
          The discount factor. The weight that the algorithm will consider future rewards. gamma = 0 is entirely greedy.
          Values higher than 1 can cause the algorithm to diverge.[1]
      reward : int
          The reward in the transition matrix
      next_state : int
          The next state of the transition matrix.
  
      Returns
      -------
      next_q : float
          The value the Q matrix will be updated with.
  
      References
      ----------
      1. Q-learning. In: Wikipedia [Internet]. 2017 [cited 2017 May 29]. Available from: https://en.wikipedia.org/w/index.php?title=Q-learning&oldid=762556833
  
      """
      max_next_step = max(q[next_state, :])
      next_q = q[state, action] + alpha * (reward + gamma * max_next_step - q[state, action])
      return next_q
  
  
  def qlearning(grid, episodes, epsilon, alpha, gamma, updated_grid):
      """The q-learning algorithm.
  
      From a grid, determine the optimal path to a reward by discovering the Q function that describes the relationship
      between states and actions.
  
      Parameters
      ----------
      grid : ndarray of ints
          The 2D grid world.
      episodes : int
          The number of successes before terminating the algorithm.
      epsilon : float
          The probability of choosing the next action randomly. Otherwise, choose the Q-optimal next step.
      alpha : float
          The learning rate. This determines how much each step is weighted.
      gamma : float
          The discount factor. The weight that the algorithm will consider future rewards. gamma = 0 is entirely greedy.
          Values higher than 1 can cause the algorithm to diverge.[1]
      updated_grid : ndarray of ints
          A new 2D grid world. Updated grid world occurs at 1000.
  
      Returns
      -------
      q_matrix : ndarray of floats
          The transition matrix that approximates the Q function of the grid world.
      steps_to_goal_list : list of ints
          The number of steps it took each iteration to get to the reward.
  
      References
      ----------
      1. Q-learning. In: Wikipedia [Internet]. 2017 [cited 2017 May 29]. Available from: https://en.wikipedia.org/w/index.php?title=Q-learning&oldid=762556833
      """
      r_matrix = make_transition_matrix(grid)
      # Initialize Q
      q_matrix = np.zeros_like(r_matrix).astype(float)
      # For each episode
      steps = 0
      steps_to_goal_list = []
      for _ in range(episodes):
          # Initialize s
          state = index_1d(5, 3)
  
          steps_to_goal = 0
  
          # Repeat (for each step in an episode)
          while state != index_1d(0, 8):
              if steps == 1000:
                  r_matrix = make_transition_matrix(updated_grid)
              # With probability \epsilon choose random action
              action = get_epsilon_greedy_action(epsilon, q_matrix, state)
              # Take action a, observe r, s'
              reward = peek_reward(r_matrix, state, action)
              next_state = peek_next_state(r_matrix, state, action)
              # Update Q
              q_matrix[state, action] = update_q(q_matrix, state, action, alpha, gamma, reward, next_state)
              # Update state
              state = next_state
              steps += 1
              steps_to_goal += 1
  
          steps_to_goal_list.append(steps_to_goal)
      # Until s is terminal
  
      return q_matrix, steps_to_goal_list
  
  
  def get_epsilon_greedy_action(epsilon, q_matrix, state):
      """Choose the next action in the q-learning algorithm.
  
      Parameters
      ----------
      epsilon : float
          The probability of choosing the next action randomly. Otherwise, choose the Q-optimal next step.
      q_matrix : ndarray of floats
          The 2D matrix representing the Q-function.
      state : int
          The current state, represented as a row number in the transition matrix.
  
      Returns
      -------
      action : int
          The next action to be taken, represented as a column number in the transition matrix.
  
  
      Notes
      -----
      If there are multiple q-optimal actions (i.e. a tie), this function selects randomly between the optimal values.
  
      """
      if random() < epsilon:
          action = randint(0, 3)
      # With probability 1 - \epsilon choose argmax Q
      else:
          action = np.random.choice(np.where(q_matrix[state] == q_matrix[state].max())[0])
      return action
  
  
  # noinspection PyTypeChecker
  @cl.command()
  @cl.option('--episodes', default=1000, help='Number of Successful Runs.')
  @cl.option('--epsilon', '-e', default=0.05, help='Probability of choosing the next action randomly (vs. greedily).')
  @cl.option('--alpha', '-a', default=0.5, help='Learning Rate.')
  @cl.option('--gamma', '-g', default=0.95, help='Discount Factor.')
  def run_q(episodes, epsilon, alpha, gamma):
      """Command Line Interface for qlearning."""
      np.set_printoptions(suppress=True)
      print("Initializing Q Learning Algorithm...")
      q, iterations = qlearning(state_grid, episodes, epsilon=epsilon, alpha=alpha, gamma=gamma,
                                updated_grid=updated_grid)
      print(q)
      print("Average steps per iteration: {}".format(sum(iterations) / len(iterations)))
  
  
  if __name__ == '__main__':
      run_q()
  