from gymnasium import spaces, utils

from miniworld.entity import Box, Ball
from miniworld.miniworld import MiniWorldEnv
from miniworld.params import DEFAULT_PARAMS
import gymnasium as gym
import numpy as np

from miniworld.envs.maze import Maze
from miniworld.entity import Box



class MazeG(Maze):
    def __init__(self, num_rows=2, num_cols=2, **kwargs):
        Maze.__init__(self, num_rows=num_rows, num_cols=num_cols, **kwargs)



def register_envs():
    module_name = __name__

    # Iterate through global names

        # Register the environment with OpenAI Gym
    task_name = "MazeG"
    gym_id = f"MiniWorld-{task_name}-v0"
    entry_point = f"{module_name}:{task_name}"

    gym.envs.registration.register(
        id=gym_id,
        entry_point=entry_point,
    )

    print('Registered env:', gym_id)


register_envs()


class my_gen_maze(object):
    def __init__(self, render_mode="rgb_array", view = 'agent'):
        #  render_mode="human", view = 'top'
        #  render_mode="rgb_array", view = 'agent'
        self.maze = gym.make("MiniWorld-MazeG-v0", render_mode = render_mode, view = view)
        self.color = ['red', 'green', 'blue', 'yellow']
        self.obj =  ['key', 'door', 'ball']
        self.obj_list = [[i,j] for i in self.color for j in self.obj]

    def reset(self):
        self.maze.reset()
        self.maze.place_entity(self.maze.agent, pos=[5.0, 0, 5.0], dir=0.66*3.14)
        self.maze.box.pos = [5.0, 5.0, 1.0]
        size = 0.4
        self.maze.box.size = np.array([size, size, size])

        self.ball = self.maze.place_entity(Ball(color="red"))
        self.ball.pos = [5.0, 0.0, 1.0]
        size = 0.4
        self.ball.size = np.array([size, size, size])

        return self.maze.render()



    def step(self, action):
        return self.maze.step(action)


    def render(self):
        return self.maze.render()

