# https://www.geeksforgeeks.org/machine-learning/q-learning-in-python/
# Agent: entity that makes decisions and learns from interaction with the environment

import numpy as np
import matplotlib.pyplot as plt

n_states = 16          # 4x4 grid → 16 possible states
n_actions = 4          # 4 actions: left, right, up, down
goal_state = 15        # target position the agent tries to reach

# Initialize Q-table (state x action) with zeros
Q_table = np.zeros((n_states, n_actions))

learning_rate = 0.8        # how strongly new information overrides old Q-values
discount_factor = 0.95     # how much future rewards matter compared to immediate ones
exploration_prob = 0.2     # chance of choosing a random action (exploration)
epochs = 1000              # number of training episodes

def get_next_state(state, action):
    row, col = divmod(state, 4)   # convert 0–15 state index into (row, col)

    # Move left
    if action == 0 and col > 0:
        col -= 1
    # Move right
    elif action == 1 and col < 3:
        col += 1
    # Move up
    elif action == 2 and row > 0:
        row -= 1
    # Move down
    elif action == 3 and row < 3:
        row += 1

    return row * 4 + col          # convert back to 0–15 state index

for epoch in range(epochs):
    current_state = np.random.randint(0, n_states)  # start from a random state

    while True:
        # ε-greedy action selection: explore or exploit
        if np.random.rand() < exploration_prob:
            action = np.random.randint(0, n_actions)    # random action
        else:
            action = np.argmax(Q_table[current_state])  # best known action

        next_state = get_next_state(current_state, action)

        # Reward: 1 for reaching the goal, 0 otherwise
        reward = 1 if next_state == goal_state else 0

        # Q-learning update rule
        Q_table[current_state, action] += learning_rate * (
            reward + discount_factor * np.max(Q_table[next_state]) - Q_table[current_state, action]
        )

        # End episode if goal reached
        if next_state == goal_state:
            break

        current_state = next_state     # move to next state
        
# Convert max Q-values into 4x4 grid for visualization
q_values_grid = np.max(Q_table, axis=1).reshape((4, 4))

# Heatmap of learned Q-values
plt.figure(figsize=(6, 6))
plt.imshow(q_values_grid, cmap='coolwarm', interpolation='nearest')
plt.colorbar(label='Q-value')
plt.title('Learned Q-values for Each State')
plt.xticks(np.arange(4), ['0', '1', '2', '3'])
plt.yticks(np.arange(4), ['0', '1', '2', '3'])
plt.gca().invert_yaxis()
plt.grid(True)

# Display numerical Q-values on each cell
for i in range(4):
    for j in range(4):
        plt.text(j, i, f'{q_values_grid[i, j]:.2f}', ha='center', va='center', color='black')

plt.show()

print("Learned Q-table:")
print(Q_table)