# -*- coding: utf-8 -*-
"""Ex08.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1sYCTiZxq-Z1zibubnLeqBRyFc7-7mcV_
"""

import numpy as np

"""# N x N Grid world"""

class Environment:
  
  
  actions = ['left', 'right', 'up', 'down']
  
  def __init__(self, grid_size=(4,4)):
    self.grid_size = grid_size 
    self.s = np.arange(self.grid_size[0] * self.grid_size[1])
    self.v = np.zeros(self.grid_size)
    self.terminal_state_indices = [self.s[0], self.s[-1]]
    # we are considering a random policy i.e. uniform probablity for each action
  
  def isTerminalState(self, state_index):
    return state_index in self.terminal_state_indices
    
  def step(self, state_index, action_index):
    num_rows, num_cols = self.grid_size
    if state_index % num_cols == 0 and self.actions[action_index] == 'left':
      return state_index, 0
    elif (state_index + 1) % num_cols == 0 and self.actions[action_index] == 'right':
      return state_index, 0
    elif state_index + 1 < num_cols and self.actions[action_index] == 'up':
      return state_index, 0
    elif state_index + 1 > (num_rows - 1) * num_cols and self.actions[action_index] == 'down':
      return state_index, 0
    else:      
      getNextState = {
          'left': -1,
          'right': 1,
          'up': -num_cols,
          'down': num_cols
      }
      next_state = getNextState[self.actions[action_index]] + state_index
      reward = -1 
      return next_state,reward
   
  def update_vtable(self, state_index): 
    num_rows, num_cols = self.grid_size
    num_valid_actions = 0
    v_sum = 0
    for action_index in np.arange(len(self.actions)):
      next_state, reward = self.step(state_index, action_index)
      if next_state == state_index:
        #it means the agent went outside the grid so this action is not possible
        continue
        
      v_sum += reward + self.v[int(next_state/num_rows), next_state%num_rows]
      num_valid_actions += 1
      # counting the number of valid actions to get the correct expectation value
    return np.round((1/num_valid_actions) * v_sum, 1)
      
      
  
  def run_iterative_Policy_Eval(self, num_iterations):
    self.v_copy = np.zeros((num_iterations, self.v.shape[0], self.v.shape[1]))
    for i in range(num_iterations):
      for state_ind in self.s:
        v_row_index, v_col_index = int(state_ind/self.grid_size[0]) , state_ind % self.grid_size[0]
        if not self.isTerminalState(state_ind):
          self.v_copy[i, v_row_index, v_col_index] = self.update_vtable(state_ind)
      self.v = self.v_copy[i]

"""# 4 X 4 Grid Environment"""

A = Environment()

print('Num States')
print(A.s)
print('termainal States')
print(A.terminal_state_indices)
print('Num Actions')
print(A.actions)
print('initial Value function')
A.v

A.run_iterative_Policy_Eval(5)

print('Final Value function')
A.v

for i in range(A.v_copy.shape[0]):
  print('----- Value Iteration after step ',i, '-------------------')
  print(A.v_copy[i])

"""# 10 x 10 Grid Experiment"""

B = Environment(grid_size= (10,10))

print('Num States')
print(B.s)
print('termainal States')
print(B.terminal_state_indices)
print('Num Actions')
print(B.actions)
print('initial Value function')
B.v

B.run_iterative_Policy_Eval(25)

print('Final Value function')
B.v

for i in range(B.v_copy.shape[0]):
  print('----- Value Iteration after step ',i, '-------------------')
  print(B.v_copy[i])

"""# Simplified Environment Impelmentation(4X4 Grid )"""

class Environment2:
  num_columns = 4
  grid_size = num_columns * num_columns
  
  actions = ['left', 'right', 'up', 'down']
  
  def __init__(self):
    self.s = np.arange(self.grid_size)
    self.v = np.zeros(self.grid_size)
    self.pi_i = 1 / len(self.actions)
    # we are considering a random policy i.e. uniform probablity for each action
  
  def isTerminalState(self, state_index):
    return state_index == 0 or state_index == 15
    
  def step(self, state_index, action_index):
    if state_index % self.num_columns == 0 and self.actions[action_index] == 'left':
      return state_index, 0
    elif (state_index + 1) % self.num_columns == 0 and self.actions[action_index] == 'right':
      return state_index, 0
    elif state_index + 1 < self.num_columns and self.actions[action_index] == 'up':
      return state_index, 0
    elif state_index + 1 > self.grid_size - self.num_columns and self.actions[action_index] == 'down':
      return state_index, 0
    else:      
      getNextState = {
          'left': -1,
          'right': 1,
          'up': -self.num_columns,
          'down': self.num_columns
      }
      next_state = getNextState[self.actions[action_index]] + state_index
      reward = -1 
      return next_state,reward
   
  def update_vtable(self, state_index): 
    num_valid_actions = 0
    v_sum = 0
    for action_index in np.arange(len(self.actions)):
      next_state, reward = self.step(state_index, action_index)
      if next_state == state_index:
        #it means the agent went outside the grid so this action is not possible
        continue
      v_sum += reward + self.v[next_state]
      num_valid_actions += 1
      # counting the number of valid actions to get the correct expectation value
    return (1/num_valid_actions) * v_sum
      
      
  
  def run_iterative_Policy_Eval(self, num_iterations):
    v_copy = np.zeros_like(self.v)
    for i in range(num_iterations):
      for state_ind in self.s:
        if not self.isTerminalState(state_ind):
          v_copy[state_ind] = self.update_vtable(state_ind)
      self.v = v_copy

F = Environment2()

F.v

F.run_iterative_Policy_Eval(1)

F.v

F.run_iterative_Policy_Eval(1)

F.v