# Experiment 10: Q-Learning Algorithm for Grid World Environment
# Requirements: pip install numpy pandas

# ==========================================
# STEP 1: Import Libraries & Environment Setup
# ==========================================
import numpy as np
import pandas as pd

# Define Grid Dimensions
GRID_ROWS = 4
GRID_COLS = 4

# Special States & Rewards
START = (3, 0)
GOAL = (0, 3)
WALL = (1, 1)

REWARD_GOAL = 10.0
REWARD_WALL = -1.0
REWARD_DEFAULT = 0.0

# Hyperparameters
ALPHA = 0.5   # Learning rate
GAMMA = 0.9   # Discount factor
EPISODES = 2000

# Actions: 0: Up, 1: Down, 2: Left, 3: Right
ACTIONS = ['Up', 'Down', 'Left', 'Right']
ACTION_MOVES = {
    'Up': (-1, 0),
    'Down': (1, 0),
    'Left': (0, -1),
    'Right': (0, 1)
}

# ==========================================
# STEP 2: Environment Step Function
# ==========================================
def step(state, action_idx):
    if state == GOAL:
        return state, 0, True

    action_name = ACTIONS[action_idx]
    dr, dc = ACTION_MOVES[action_name]
    next_r = state[0] + dr
    next_c = state[1] + dc

    # Check boundary collision
    if next_r < 0 or next_r >= GRID_ROWS or next_c < 0 or next_c >= GRID_COLS:
        next_state = state
    else:
        next_state = (next_r, next_c)

    # Determine reward and terminal condition
    if next_state == WALL:
        # Wall penalty; agent hits wall
        reward = REWARD_WALL
        is_terminal = False
    elif next_state == GOAL:
        reward = REWARD_GOAL
        is_terminal = True
    else:
        reward = REWARD_DEFAULT
        is_terminal = False

    return next_state, reward, is_terminal

# ==========================================
# STEP 3: Q-Learning Algorithm Implementation
# ==========================================
def run_q_learning():
    # Initialize Q-table with zeros: shape (4, 4, 4)
    q_table = np.zeros((GRID_ROWS, GRID_COLS, len(ACTIONS)))

    for episode in range(EPISODES):
        state = START
        is_terminal = False

        while not is_terminal:
            # Epsilon-greedy action selection
            epsilon = max(0.1, 1.0 - episode / 1000)
            if np.random.rand() < epsilon:
                action_idx = np.random.choice(len(ACTIONS))
            else:
                action_idx = np.argmax(q_table[state[0], state[1]])

            next_state, reward, is_terminal = step(state, action_idx)

            # Target Q-value calculation
            if is_terminal:
                max_next_q = 0.0
            else:
                max_next_q = np.max(q_table[next_state[0], next_state[1]])

            td_target = reward + GAMMA * max_next_q
            td_error = td_target - q_table[state[0], state[1], action_idx]

            # Q-value update rule
            q_table[state[0], state[1], action_idx] += ALPHA * td_error

            state = next_state

    return q_table

# Run training
q_table = run_q_learning()

# ==========================================
# STEP 4: Formatted Q-Table Display
# ==========================================
rows = []
for r in range(GRID_ROWS):
    for c in range(GRID_COLS):
        state = (r, c)
        if state == GOAL:
            rows.append({"State (r,c)": f"({r},{c}) [Goal]", "Up": "-", "Down": "-", "Left": "-", "Right": "-", "Optimal Action": "Terminal"})
        elif state == WALL:
            rows.append({"State (r,c)": f"({r},{c}) [Wall]", "Up": "-", "Down": "-", "Left": "-", "Right": "-", "Optimal Action": "Obstacle"})
        else:
            q_vals = q_table[r, c]
            best_a = ACTIONS[np.argmax(q_vals)]
            rows.append({
                "State (r,c)": f"({r},{c})",
                "Up": f"{q_vals[0]:.4f}",
                "Down": f"{q_vals[1]:.4f}",
                "Left": f"{q_vals[2]:.4f}",
                "Right": f"{q_vals[3]:.4f}",
                "Optimal Action": best_a
            })

df_qtable = pd.DataFrame(rows)
print("--- CONVERGED Q-TABLE ---")
print(df_qtable.to_string(index=False))
