text: # -*- coding: utf-8 -*-
  """
  Created on Tue Dec 24 20:43:09 2019
  
  @author: elif.ayvali
  """
  
  import sys
  import gym
  import pandas as pd
  import numpy as np
  import matplotlib.collections as mc
  import matplotlib.pyplot as plt
  from helpers import create_uniform_grid,discretize,visualize_samples
  
  class QLearningAgent:
      """Q-Learning agent that can act on a continuous state space by discretizing it."""
  
      def __init__(self, env, state_grid, alpha=0.02, gamma=0.99, epsilon=1.0, epsilon_decay_rate=0.9995, min_epsilon=.01, seed=505):
          """Initialize variables, create grid for discretization."""
          # Environment info
          self.env = env
          self.state_grid = state_grid
          self.state_size = tuple(len(splits) + 1 for splits in self.state_grid)  # n-dimensional state space
          self.action_size = self.env.action_space.n  # 1-dimensional discrete action space
          self.seed = np.random.seed(seed)
          print("Environment:", self.env)
          print("State space size:", self.state_size)
          print("Action space size:", self.action_size)
          
          # Learning parameters
          self.alpha = alpha  # learning rate
          self.gamma = gamma  # discount factor
          self.epsilon = self.initial_epsilon = epsilon  # initial exploration rate
          self.epsilon_decay_rate = epsilon_decay_rate # how quickly should we decrease epsilon
          self.min_epsilon = min_epsilon
          
          # Create Q-table
          self.q_table = np.zeros(shape=(self.state_size + (self.action_size,)))
          print("Q table size:", self.q_table.shape)
  
      def preprocess_state(self, state):
          """Map a continuous state to its discretized representation."""
          return tuple(discretize(state,self.state_grid))
  
      def reset_episode(self, state):
          """Reset variables for a new episode."""
          # Gradually decrease exploration rate
          self.epsilon *= self.epsilon_decay_rate
          self.epsilon = max(self.epsilon, self.min_epsilon)
  
          # Decide initial action
          self.last_state = self.preprocess_state(state)
          self.last_action = np.argmax(self.q_table[self.last_state])
          return self.last_action
      
      def reset_exploration(self, epsilon=None):
          """Reset exploration rate used when training."""
          self.epsilon = epsilon if epsilon is not None else self.initial_epsilon
  
      def act(self, state, reward=None, done=None, mode='train'):
          """Pick next action and update internal Q table (when mode != 'test')."""
          state = self.preprocess_state(state)
          if mode == 'test':
              # Test mode: Simply produce an action
              action = np.argmax(self.q_table[state])
          else:
              # Train mode (default): Update Q table, pick next action
              # Note: We update the Q table entry for the *last* (state, action) pair with current state, reward
              self.q_table[self.last_state + (self.last_action,)] += self.alpha * \
                  (reward + self.gamma * max(self.q_table[state]) - self.q_table[self.last_state + (self.last_action,)])
  
              # Exploration vs. exploitation
              do_exploration = np.random.uniform(0, 1) < self.epsilon
              if do_exploration:
                  # Pick a random action
                  action = np.random.randint(0, self.action_size)
              else:
                  # Pick the best action from Q table
                  action = np.argmax(self.q_table[state])
  
          # Roll over current state, action for next step
          self.last_state = state
          self.last_action = action
          return action
  
  def run(agent, env, num_episodes=20000, mode='train'):
      """Run agent in given reinforcement learning environment and return scores."""
      scores = []
      max_avg_score = -np.inf
      for i_episode in range(1, num_episodes+1):
          # Initialize episode
          state = env.reset()
          action = agent.reset_episode(state)
          total_reward = 0
          done = False
  
          # Roll out steps until done
          while not done:
              state, reward, done, info = env.step(action)
              total_reward += reward
              action = agent.act(state, reward, done, mode)
  
          # Save final score
          scores.append(total_reward)
          
          # Print episode stats
          if mode == 'train':
              if len(scores) > 100:
                  avg_score = np.mean(scores[-100:])
                  if avg_score > max_avg_score:
                      max_avg_score = avg_score
              if i_episode % 100 == 0:
                  print("\rEpisode {}/{} | Max Average Score: {}".format(i_episode, num_episodes, max_avg_score), end="")
                  sys.stdout.flush()
  
      return scores
  
  def plot_scores(scores, rolling_window=100):
      """Plot scores and optional rolling mean using specified window."""
      plt.plot(scores); plt.title("Scores");
      rolling_mean = pd.Series(scores).rolling(rolling_window).mean()
      plt.plot(rolling_mean);
      return rolling_mean
  
  def plot_q_table(q_table):
      """Visualize max Q-value for each state and corresponding action.
      Q-table is of size MxNxA, 
      where (M, N) is the size of the state space, 
      and A is the size of the action space.
      We are interested in the maximum Q-value for each state,
      and the corresponding (best) action associated with that value."""
      q_image = np.max(q_table, axis=2)       # max Q-value for each state
      q_actions = np.argmax(q_table, axis=2)  # best action for each state
  
      fig, ax = plt.subplots(figsize=(10, 10))
      cax = ax.imshow(q_image, cmap='jet');
      cbar = fig.colorbar(cax)
      for x in range(q_image.shape[0]):
          for y in range(q_image.shape[1]):
              ax.text(x, y, q_actions[x, y], color='white',
                      horizontalalignment='center', verticalalignment='center')
      ax.grid(False)
      ax.set_title("Q-table, size: {}".format(q_table.shape))
      ax.set_xlabel('position')
      ax.set_ylabel('velocity')
      
  # Create an environment 
  env = gym.make('MountainCar-v0')
  # Create a grid to discretize the state space of the RL problem
  low = [-1.0, -5.0]
  high = [1.0, 5.0]
  state_grid = create_uniform_grid(env.observation_space.low, env.observation_space.high, bins=(10, 10))    
  q_agent = QLearningAgent(env, state_grid)
  scores = run(q_agent, env)
  # Plot scores obtained per episode
  plt.plot(scores); plt.title("Scores")
  #If the scores are noisy, it might be difficult to tell whether your agent is actually learning.
  # To find the underlying trend, plot a rolling mean of the scores.
  rolling_mean = plot_scores(scores)
  # Run in test mode and analyze scores obtained
  test_scores = run(q_agent, env, num_episodes=100, mode='test')
  print("[TEST] Completed {} episodes with avg. score = {}".format(len(test_scores), np.mean(test_scores)))
  _ = plot_scores(test_scores)
  plot_q_table(q_agent.q_table)
  
  #Improve the agent by increasing state space
  state_grid_new = create_uniform_grid(env.observation_space.low, env.observation_space.high, bins=(20, 20))
  q_agent_new = QLearningAgent(env, state_grid_new)
  q_agent_new.scores = []  # initialize a list to store scores for this agent
  #This line can be run multiple times, and scores will get accumulated
  q_agent_new.scores += run(q_agent_new, env, num_episodes=50000)  # accumulate scores
  rolling_mean_new = plot_scores(q_agent_new.scores)
  
  # Run in test mode and analyze scores obtained
  test_scores = run(q_agent_new, env, num_episodes=100, mode='test')
  print("[TEST] Completed {} episodes with avg. score = {}".format(len(test_scores), np.mean(test_scores)))
  _ = plot_scores(test_scores)
  # Visualize the learned Q-table
  plot_q_table(q_agent_new.q_table)
  
  