"""A custom observation_function for retro.examples.snake.

Crops the board to a window centered on the snake's head before encoding,
using egocentric_board()/encode_board() from retro_gamer.observation as
building blocks. This replaces retro_gamer's old built-in egocentric flags —
cropping is now just something an observation_function does for itself.

Point a run's config.toml at this with:

    [metadata]
    observation_function = "snake_observation:egocentric_observation"
    board_size = [17, 17]   # must match RADIUS below: 2*8+1 = 17
"""
import numpy as np
from retro.views.headless import HeadlessView
from retro_gamer.observation import egocentric_board, encode_board, encode_state

CHARACTER_SET = ["@", "*", ">", "<", "^", "v"]
RADIUS = 8


def egocentric_observation(game):
    """Board cropped to a (2*RADIUS+1)^2 window around the snake's head, plus apple_dx/apple_dy."""
    view = HeadlessView()
    view.on_game_start(game)
    view.render(game)
    head = game.get_agent_by_name("Snake head")
    cropped = egocentric_board(view.board_characters, head.position, RADIUS)
    board_vec = encode_board(cropped, CHARACTER_SET).flatten()
    extras = encode_state(game.state, ["apple_dx", "apple_dy"])
    return np.concatenate([board_vec, extras])