# A simple simulated world for Rosie the Robot, based on Chapter 8 of
# "Artificial Intelligence: A Guide for Thinking Humans" by Melanie Mitchell

# version 1: simulates one episode of Rosie performing random actions

#---------------------------------------------------------------------------
import random

class RosieWorld:

    def __init__(self, size=26):
        self.world_size = size
        self.robot_location = random.randrange(size)
        self.ball_location = 0

    def __str__(self):
        s = ""
        for i in range(self.world_size):
            if self.robot_location == i and self.ball_location == i:
                s += "oR "
            elif self.ball_location == i:
                s += "o  "
            elif self.robot_location == i:
                s += "R  "
            else:
                s += "_  "
        return s

    def forward(self):
        if self.robot_location > 0:
            self.robot_location -= 1
        return 0

    def backward(self):
        if self.robot_location < self.world_size-1:
            self.robot_location += 1
        return 0

    def kick(self):
        if self.robot_location == self.ball_location:
            self.ball_location = None
            return 10
        else:
            return 0

    def can_see_ball(self):
        if self.ball_location == None:
            return False
        else:
            return True

#---------------------------------------------------------------------------
# simulates one episode of Rosie performing actions in her world

ACTION_NAMES = ['forward', 'back', 'kick']
ACTION_CODES = [0, 1, 2]

def run_episode():
    rosie = RosieWorld()
    print(f"initial state:\t{rosie}")
    steps = 0
    while rosie.can_see_ball():
        # choose an action randomly
        action = random.choice(ACTION_CODES)
        # perform action and observe reward
        if action == 0:
            reward = rosie.forward()
        elif action == 1:
            reward = rosie.backward()
        elif action == 2:
            reward = rosie.kick()
        print(f"action: {ACTION_NAMES[action]}\t{rosie}")
        steps += 1
    plural = '' if steps == 1 else 's'
    print(f"Rosie kicked the ball after {steps} step{plural}!")
    return steps

#---------------------------------------------------------------------------