"""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.
"""

from q_learning import QLearning
import babysnake
from retro.input import ProgrammaticInput
from retro_gamer import GameEnvironment, GameMetadata

def train():
    trainer = QLearning()
    env = GameEnvironment(babysnake.create_game, GameMetadata.from_pyproject("babysnake"))
    return trainer.train(env, babysnake.ACTIONS)

def watch(Q):
    inp = ProgrammaticInput()

    class PolicyInput:
        """An input source that picks actions from the Q-table."""
        def collect(self):
            state = babysnake.get_state(game)
            q, action = sorted([(Q.get((state, a), 0), a) for a in babysnake.ACTIONS], reverse=True)[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 = train()
    print(f"\nDone. Q-table has {len(Q)} entries.")
    print("\nWatching trained agent (press Enter or Escape to quit)...")
    watch(Q)