"""BabySnake: a 4×4 grid game where an agent collects food.

State: (agent_x, agent_y, food_x, food_y) — four integers.
The agent starts with 50 energy. Each step costs 1 energy.
Collecting food restores 30 energy and adds 1 to the score.
The game ends when energy reaches 0.
"""

from random import randint, choice
from retro.game import Game

BOARD_SIZE = 4
START_ENERGY = 50
FOOD_ENERGY = 30

class Forager:
    """The player agent. Controlled with arrow keys."""
    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['energy'] -= 1
        game.state['reward'] -= 0.01

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

        ax, ay = self.position
        fx, fy = game.get_agent_by_name("Food").position
        game.state['agent_x'] = ax
        game.state['agent_y'] = ay
        game.state['food_x'] = fx
        game.state['food_y'] = fy

        if game.state['energy'] <= 0:
            game.end()


class Food:
    """The food item. Respawns at a random empty position 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 BabySnake game."""
    forager = Forager()
    food = Food()
    bw = bh = BOARD_SIZE
    game = Game(
        [forager, food],
        {
            'score': 0,
            'reward': 0.0,
            'energy': START_ENERGY,
            'agent_x': 0,
            'agent_y': 0,
            'food_x': 0,
            'food_y': 0,
        },
        board_size=(bw, bh),
        framerate=6,
    )
    forager.position = (randint(0, bw - 1), randint(0, bh - 1))
    food.relocate(game)
    ax, ay = forager.position
    fx, fy = food.position
    game.state['agent_x'] = ax
    game.state['agent_y'] = ay
    game.state['food_x'] = fx
    game.state['food_y'] = fy
    return game


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