"""Frogger: guide a frog across busy traffic lanes to reach the far side.

The frog starts at the bottom row. Cars move across the lanes between the
start and goal. Each lane has cars moving at different speeds in alternating
directions. The frog earns +10 for each row advanced, +50 for reaching the
top row, and -10 for being hit by a car or falling off the edge. Episodes end
when the frog reaches the top, gets hit, or energy runs out.
"""

from random import randint
from retro.game import Game

BOARD_WIDTH = 20
BOARD_HEIGHT = 12
NUM_LANES = BOARD_HEIGHT - 2
START_ENERGY = 200


class Frog:
    name = "Frog"
    character = 'O'
    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.UP

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

    def play_turn(self, game):
        bw, bh = game.board_size
        x, y = self.position
        dx, dy = self._direction
        nx, ny = x + dx, y + dy

        if not (0 <= nx < bw):
            game.state['reward'] -= 10
            game.state['energy'] -= 50
            self._reset(game)
            return

        if not (0 <= ny < bh):
            if ny < 0:
                game.state['score'] += 50
                game.state['reward'] += 50
            else:
                game.state['reward'] -= 5
            self._reset(game)
            return

        prev_y = y
        self.position = (nx, ny)
        game.state['energy'] -= 1
        game.state['reward'] -= 0.01

        if ny < prev_y:
            advancement = prev_y - ny
            game.state['score'] += advancement * 10
            game.state['reward'] += advancement * 5

        for agent in game.agents:
            if hasattr(agent, '_is_car') and agent.position == self.position:
                game.state['reward'] -= 10
                game.state['energy'] -= 50
                self._reset(game)
                return

        bw, bh = game.board_size
        fx, fy = self.position
        game.state['frog_x'] = fx / bw
        game.state['frog_y'] = fy / bh

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

    def _reset(self, game):
        bw, bh = game.board_size
        self.position = (bw // 2, bh - 1)
        self._direction = self.UP
        if game.state['energy'] <= 0:
            game.end()


class Car:
    _is_car = True
    character = 'X'
    color = "red_on_black"

    def __init__(self, lane, speed, direction, start_x, board_width):
        self.name = f"Car {lane}_{start_x}"
        self._lane = lane
        self._speed = speed
        self._direction = direction
        self._board_width = board_width
        self._step = 0
        self.position = (start_x, lane)

    def play_turn(self, game):
        self._step += 1
        if self._step < self._speed:
            return
        self._step = 0
        x, y = self.position
        x = (x + self._direction) % self._board_width
        self.position = (x, y)

        frog = game.get_agent_by_name("Frog")
        if frog.position == self.position:
            frog._reset(game)
            game.state['reward'] -= 10
            game.state['energy'] -= 50


def create_game():
    bw, bh = BOARD_WIDTH, BOARD_HEIGHT
    frog = Frog()
    frog.position = (bw // 2, bh - 1)

    agents = [frog]
    car_id = 0
    for lane_idx, row in enumerate(range(1, bh - 1)):
        direction = 1 if lane_idx % 2 == 0 else -1
        speed = 2 + (lane_idx % 3)
        num_cars = 2 + (lane_idx % 3)
        spacing = bw // num_cars
        for i in range(num_cars):
            start_x = (i * spacing + lane_idx * 3) % bw
            agents.append(Car(row, speed, direction, start_x, bw))
            car_id += 1

    game = Game(
        agents,
        {
            'score': 0,
            'reward': 0.0,
            'energy': START_ENERGY,
            'frog_x': (bw // 2) / bw,
            'frog_y': (bh - 1) / bh,
        },
        board_size=(bw, bh),
        framerate=8,
        show_state=['score', 'energy'],
    )
    return game


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