import numpy as np
import gym
import utils as util
import maze as m

class MazeEnv(gym.Env):
    def __init__(self):
        super(MazeEnv, self).__init__()
    
    @property
    def observation(self):
        resize_observation = util.resize_nearest( self._maze.observation, 42, 42)
        return [resize_observation, resize_observation, resize_observation]
    
    @property
    def posx(self):
        return self._maze.posx

    @property
    def posy(self):
        return self._maze.posy

    @property
    def maze_size(self):
        return self._maze.maze_size
    
    def load_maze(self, maze, repeat):
        self._maze = maze
        self._repeat = repeat
        self.action_space = gym.spaces.Discrete(self._maze.action_num)
        self.observation_space = gym.spaces.Box(low=0, high=self._maze.maze_size-1, shape=(2,), dtype=np.float32)

    def reset(self):
        self._maze.reset()
        return self.observation

    def step(self, action):
        t_reward = 0.0
        done = False
        info = {}
        for _ in range(self._repeat):
            reward, done, = self._maze.transit(action)
            t_reward += reward
            if done:
                break
        return self.observation, t_reward, done, info

    def render(self, mode='console', close=False):
        if mode != 'console':
            raise NotImplementedError()
        self._maze.print_maze()