"""Q-learning agent.

This file contains starter code for implementing a Q-learning agent.
You need to fill in two functions:
  - choose_action: select an action using an epsilon-greedy policy
  - update_q: update the Q-table using the Bellman equation

Both functions are written here with no reference to BabySnake or any
particular game — they work on any (state, action) Q-table, given a list of
possible actions. train_babysnake.py is the module that wires these functions
up to BabySnake specifically.

Before training, get your implementation passing the tests in
test_q_learning.py:

    python test_q_learning.py

Errors caught here are much easier to track down than errors discovered
during training. Once the tests pass, run:

    python train_babysnake.py
"""

import random


def choose_action(q_table, state, actions, epsilon):
    """Choose an action using an epsilon-greedy policy.

    With probability `epsilon`, return a random action from `actions`.
    Otherwise, return the action with the highest Q-value in `q_table`
    for the given `state`. If a (state, action) pair has not been seen
    before, treat its Q-value as 0.

    Arguments:
        q_table (dict): Maps (state, action) -> Q-value.
        state: The current state.
        actions (list): The actions available to choose from.
        epsilon (float): Exploration rate, between 0.0 and 1.0.

    Returns:
        One of the values in `actions`.

    Hint: random.random() returns a float in [0.0, 1.0).
          random.choice(actions) returns a random action.
          q_table.get(key, default) is handy for missing entries.
    """
    raise NotImplementedError("Fill in choose_action")


def update_q(q_table, state, action, reward, next_state, actions, alpha, gamma):
    """Update one entry of the Q-table using the Bellman equation.

    The update rule is:

        Q(s, a) <- Q(s, a) + alpha * (r + gamma * max_a' Q(s', a') - Q(s, a))

    where:
        s, a       — the state we were in and the action we took
        r          — the reward we received
        s'         — the state we ended up in
        max_a' ... — the best possible Q-value from the new state

    Arguments:
        q_table (dict): Maps (state, action) -> Q-value (modified in place).
        state: The state before the action.
        action: The action taken.
        reward (float): The reward received.
        next_state: The state after the action.
        actions (list): The actions available from `next_state`.
        alpha (float): Learning rate (how much to update).
        gamma (float): Discount factor (how much to value future rewards).

    Returns:
        None — modifies q_table in place.

    Hint: Q-values for unseen (state, action) pairs start at 0.
    """
    raise NotImplementedError("Fill in update_q")