import gym
from gym import spaces

from deepmind import PillEater, observation_as_rgb
import numpy as np

class MiniPacman:
    def __init__(self, mode, frame_cap):
        self.mode      = mode
        self.frame_cap = frame_cap

        self.env = PillEater(mode=mode, frame_cap=frame_cap)

        self.action_space      = spaces.Discrete(5)
#        self.observation_space = spaces.Box(low=0, high=1.0, shape=(3, 15, 19))
        # adapted for smaller board
        self.observation_space = spaces.Box(low=0, high=1.0, shape=(3, 10, 7))

    def step(self, action):
        self.env.step(action)
        env_reward, env_pcontinue, env_frame = self.env.observation()
        self.done = env_pcontinue != 1
        env_frame = env_frame.transpose(2, 0, 1)
        return env_frame, env_reward, self.done, {}

    # added option to give pacman/ghost/pill initial position
    def reset(self, init_pos=False, pos_pillman=np.array([0,0]), pos_ghost=np.array([0,0]), dir_ghost=0, pos_pill=np.array([0,0])):
        image, _, _ = self.env.start(init_pos, pos_pillman, pos_ghost, dir_ghost, pos_pill)
        image = observation_as_rgb(image)
        self.done = False
        image = image.transpose(2, 0, 1)
        return image