"""Forager: an 8×8 grid game where an agent collects food.

The agent moves in four directions collecting food that respawns on collection.
The game runs indefinitely; retro-gamer's max_turns_per_episode controls episode length.

Observation features for retro-gamer:
  food_dx: (food_x - agent_x) / board_width  (positive = food is to the right)
  food_dy: (food_y - agent_y) / board_height (positive = food is below)
"""

from random import randint
from retro.game import Game

BOARD_SIZE = 8


class Forager:
    """The player agent."""
    name = "Forager"
    character = '@'
    color = "green_on_black"
    position = (0, 0)

    UP = (0, -1)
    DOWN = (0, 1)
    LEFT = (-1, 0)
    RIGHT = (1, 0)

    def __init__(self):
        self._direction = self.RIGHT

    def handle_keystroke(self, keystroke, game):
        if keystroke.name == "KEY_RIGHT":
            self._direction = self.RIGHT
        elif keystroke.name == "KEY_UP":
            self._direction = self.UP
        elif keystroke.name == "KEY_LEFT":
            self._direction = self.LEFT
        elif keystroke.name == "KEY_DOWN":
            self._direction = self.DOWN

    def play_turn(self, game):
        x, y = self.position
        dx, dy = self._direction
        new_pos = (x + dx, y + dy)
        if game.on_board(new_pos):
            self.position = new_pos

        game.state['reward'] -= 0.01

        food = game.get_agent_by_name("Food")
        if self.position == food.position:
            food.relocate(game)
            game.state['score'] += 1
            game.state['reward'] += 1.0

        bw, bh = game.board_size
        ax, ay = self.position
        fx, fy = game.get_agent_by_name("Food").position
        game.state['food_dx'] = (fx - ax) / bw
        game.state['food_dy'] = (fy - ay) / bh


class Food:
    """The food item. Respawns when collected."""
    name = "Food"
    character = '*'
    color = "yellow_on_black"
    position = (0, 0)

    def relocate(self, game):
        bw, bh = game.board_size
        forager = game.get_agent_by_name("Forager")
        while True:
            pos = (randint(0, bw - 1), randint(0, bh - 1))
            if pos != forager.position:
                self.position = pos
                return


def create_game():
    """Return a fresh Forager game."""
    forager = Forager()
    food = Food()
    bw = bh = BOARD_SIZE
    game = Game(
        [forager, food],
        {'score': 0, 'reward': 0.0, 'food_dx': 0.0, 'food_dy': 0.0},
        board_size=(bw, bh),
        framerate=12,
    )
    forager.position = (randint(0, bw - 1), randint(0, bh - 1))
    food.relocate(game)
    bw, bh = game.board_size
    ax, ay = forager.position
    fx, fy = food.position
    game.state['food_dx'] = (fx - ax) / bw
    game.state['food_dy'] = (fy - ay) / bh
    return game


if __name__ == '__main__':
    create_game().play()