"""Train a Q-learning agent to play BabySnake, then watch it play.

Run this file to train and watch:

    python train_babysnake.py

This module wires the generic Q-learning algorithm in q_learning.py up to
BabySnake specifically, using retro_gamer.GameEnvironment (configured by
babysnake/pyproject.toml's observation_function) as the environment.
"""

import babysnake
import q_learning
from babysnake_env import ACTIONS, get_state
from retro.input import ProgrammaticInput
from retro_gamer import GameEnvironment, GameMetadata


def train(
    episodes=1000,
    alpha=0.1,
    gamma=0.95,
    epsilon=1.0,
    epsilon_decay=0.995,
    epsilon_min=0.05,
    max_steps_per_episode=500,
):
    """Train a Q-learning agent on BabySnake.

    Arguments:
        episodes (int): How many episodes to run.
        alpha (float): Learning rate.
        gamma (float): Discount factor.
        epsilon (float): Starting exploration rate.
        epsilon_decay (float): Multiply epsilon by this each episode.
        epsilon_min (float): Epsilon never falls below this.
        max_steps_per_episode (int): Safety cutoff. Without this, a lucky
            random walk that keeps finding food (each pickup restores more
            energy than a turn costs) can make an episode run far longer
            than intended, or even effectively forever.

    Returns:
        dict: The trained Q-table.
    """
    q_table = {}
    env = GameEnvironment(babysnake.create_game, GameMetadata.from_pyproject("babysnake"))

    for episode in range(episodes):
        state = env.reset()
        total_reward = 0.0

        for _ in range(max_steps_per_episode):
            if not env.game.playing:
                break
            action = q_learning.choose_action(q_table, state, ACTIONS, epsilon)
            next_state, reward, done = env.step(action)
            q_learning.update_q(q_table, state, action, reward, next_state, ACTIONS, alpha, gamma)
            state = next_state
            total_reward += reward

        epsilon = max(epsilon_min, epsilon * epsilon_decay)

        if (episode + 1) % 100 == 0:
            print(
                f"Episode {episode + 1:5d}  "
                f"reward={total_reward:6.1f}  "
                f"score={env.game.state['score']}  "
                f"epsilon={epsilon:.3f}  "
                f"q_entries={len(q_table)}"
            )

    return q_table


def watch(q_table=None):
    """Watch the trained agent play in the terminal.

    Arguments:
        q_table (dict | None): A trained Q-table. If None, trains first.
    """
    if q_table is None:
        print("Training first...")
        q_table = train()

    inp = ProgrammaticInput()

    class PolicyInput:
        """An input source that picks actions from the Q-table."""
        def collect(self):
            state = get_state(game)
            action = q_learning.choose_action(q_table, state, ACTIONS, epsilon=0.0)
            inp.press(action)
            return inp.collect()

    game = babysnake.create_game()
    game.play(input_source=PolicyInput())


if __name__ == '__main__':
    print("Training Q-learning agent on BabySnake...")
    q_table = train()
    print(f"\nDone. Q-table has {len(q_table)} entries.")
    print("\nWatching trained agent (press Enter or Escape to quit)...")
    watch(q_table)