"""Solution for q_learning.py — remove before publishing to students."""

import random


def choose_action(q_table, state, actions, epsilon):
    if random.random() < epsilon:
        return random.choice(actions)
    q_values = [q_table.get((state, a), 0.0) for a in actions]
    return actions[q_values.index(max(q_values))]


def update_q(q_table, state, action, reward, next_state, actions, alpha, gamma):
    old_q = q_table.get((state, action), 0.0)
    next_q_values = [q_table.get((next_state, a), 0.0) for a in actions]
    best_next_q = max(next_q_values)
    new_q = old_q + alpha * (reward + gamma * best_next_q - old_q)
    q_table[(state, action)] = new_q