from gym_minigrid.minigrid import *
from gym_minigrid.register import register
import random

class FourRoomsEnv(MiniGridEnv):
    """
    Classic 4 rooms gridworld environment.
    Can specify agent and goal position, if not it set at random.
    """

    def __init__(self, agent_pos=None, goal_pos=None):
        self._agent_default_pos = agent_pos
        self._goal_default_pos = goal_pos
        super().__init__(grid_size=15, max_steps=1000)

    def _gen_grid(self, width, height):
        # Create the grid
        self.grid = Grid(width, height)

        # Generate the surrounding walls
        self.grid.horz_wall(0, 0)
        self.grid.horz_wall(0, height - 1)
        self.grid.vert_wall(0, 0)
        self.grid.vert_wall(width - 1, 0)

        room_w = width // 2
        room_h = height // 2

        # For each row of rooms
        for j in range(0, 2):

            # For each column
            for i in range(0, 2):
                xL = i * room_w
                yT = j * room_h
                xR = xL + room_w
                yB = yT + room_h

                # Bottom wall and door
                if i + 1 < 2:
                    self.grid.vert_wall(xR, yT, room_h)
                    pos = (xR, self._rand_int(yT + 1, yB))
                    self.grid.set(*pos, None)

                # Bottom wall and door
                if j + 1 < 2:
                    self.grid.horz_wall(xL, yB, room_w)
                    pos = (self._rand_int(xL + 1, xR), yB)
                    self.grid.set(*pos, None)

        # Randomize the player start position and orientation
        if self._agent_default_pos is not None:
            self.agent_pos = self._agent_default_pos
            self.grid.set(*self._agent_default_pos, None)
            self.agent_dir = self._rand_int(0, 4)  # assuming random start direction
        else:
            self.place_agent()

        if self._goal_default_pos is not None:
            goal = Goal()
            self.put_obj(goal, *self._goal_default_pos)
            goal.init_pos, goal.cur_pos = self._goal_default_pos
        else:
            self.place_obj(Goal())

        self.mission = 'Reach the goal'

    def step(self, action):
        obs, reward, done, info = MiniGridEnv.step(self, action)
        return obs, reward, done, info

class NoRoomsDetEnv(MiniGridEnv):
    """
    Classic 4 rooms gridworld environment.
    Can specify agent and goal position, if not it set at random.
    """

    def __init__(self, agent_pos=None, goal_pos=None):
        self._agent_default_pos = agent_pos
        self._goal_default_pos = goal_pos
        super().__init__(grid_size=19, max_steps=1000)

    def _gen_grid(self, width, height):
        # Create the grid
        self.grid = Grid(width, height)

        # Generate the surrounding walls
        self.grid.horz_wall(0, 0) # obj_type=Lava)
        self.grid.horz_wall(0, height - 1)
        self.grid.vert_wall(0, 0)
        self.grid.vert_wall(width - 1, 0)

        # Randomize the player start position and orientation
        if self._agent_default_pos is not None:
            self.agent_pos = self._agent_default_pos
            self.grid.set(*self._agent_default_pos, None)
            self.agent_dir = self._rand_int(0, 4)  # assuming random start direction
        else:
            self.place_agent()

        if self._goal_default_pos is not None:
            goal = Goal()
            self.put_obj(goal, *self._goal_default_pos)
            goal.init_pos, goal.cur_pos = self._goal_default_pos

        self.mission = 'Reach the goal'

    def step(self, action):
        obs, reward, done, info = MiniGridEnv.step(self, action)
        return obs, reward, done, info


class SmallNoRoomsDetEnv(MiniGridEnv):
    """
    Classic 4 rooms gridworld environment.
    Can specify agent and goal position, if not it set at random.
    """

    def __init__(self, agent_pos=None, goal_pos=None):
        self._agent_default_pos = agent_pos
        self._goal_default_pos = goal_pos
        super().__init__(grid_size=8, max_steps=1000)

    def _gen_grid(self, width, height):
        # Create the grid
        self.grid = Grid(width, height)

        # Generate the surrounding walls
        self.grid.horz_wall(0, 0) # obj_type=Lava)
        self.grid.horz_wall(0, height - 1)
        self.grid.vert_wall(0, 0)
        self.grid.vert_wall(width - 1, 0)

        # Randomize the player start position and orientation
        if self._agent_default_pos is not None:
            self.agent_pos = self._agent_default_pos
            self.grid.set(*self._agent_default_pos, None)
            self.agent_dir = self._rand_int(0, 4)  # assuming random start direction
        else:
            self.place_agent()

        if self._goal_default_pos is not None:
            goal = Goal()
            self.put_obj(goal, *self._goal_default_pos)
            goal.init_pos, goal.cur_pos = self._goal_default_pos
        else:
            self.place_obj(Goal())

        self.mission = 'Reach the goal'

    def step(self, action):
        obs, reward, done, info = MiniGridEnv.step(self, action)
        return obs, reward, done, info



class FourRoomsDetEnv(MiniGridEnv):
    """
    Classic 4 rooms gridworld environment.
    Can specify agent and goal position, if not it set at random.
    """

    def __init__(self, agent_pos=None, goal_pos=None):
        self._agent_default_pos = agent_pos
        self._goal_default_pos = goal_pos
        super().__init__(grid_size=9, max_steps=1000)

    def _gen_grid(self, width, height):
        # Create the grid
        self.grid = Grid(width, height)

        # Generate the surrounding walls
        self.grid.horz_wall(0, 0) #obj_type=Lava)
        self.grid.horz_wall(0, height - 1)
        self.grid.vert_wall(0, 0)
        self.grid.vert_wall(width - 1, 0)

        # 4 rooms
        self.grid.horz_wall(0, 4)
        self.grid.vert_wall(4, 0)

        # gates
        self.grid.set(2, 4, None)
        # self.grid.set(4, 7, None)
        self.grid.set(4, 2, None)
        # self.grid.set(7, 4, None)
        # self.grid.set(7, 10, None)
        self.grid.set(4, 6, None)
        # self.grid.set(10, 7, None)
        self.grid.set(6, 4, None)


        # Randomize the player start position and orientation
        if self._agent_default_pos is not None:
            self.agent_pos = self._agent_default_pos
            self.grid.set(*self._agent_default_pos, None)
            self.agent_dir = self._rand_int(0, 4)  # assuming random start direction
        else:
            self.place_agent()

        if self._goal_default_pos is not None:
            goal = Goal()
            self.put_obj(goal, *self._goal_default_pos)
            goal.init_pos, goal.cur_pos = self._goal_default_pos

        self.mission = 'Reach the goal'

    def step(self, action):
        obs, reward, done, info = MiniGridEnv.step(self, action)
        return obs, reward, done, info



class NineRoomsDetEnv(MiniGridEnv):
    """
    Classic 4 rooms gridworld environment.
    Can specify agent and goal position, if not it set at random.
    """

    def __init__(self, agent_pos=None, goal_pos=None):
        self._agent_default_pos = agent_pos
        self._goal_default_pos = goal_pos
        super().__init__(grid_size=19, max_steps=1000)    # always neet to set this up to more then the wraper one

    def _gen_grid(self, width, height):
        # Create the grid
        self.grid = Grid(width, height)

        # Generate the surrounding walls
        self.grid.horz_wall(0, 0)
        self.grid.horz_wall(0, height - 1)
        self.grid.vert_wall(0, 0)
        self.grid.vert_wall(width - 1, 0)

        # 9 rooms
        self.grid.horz_wall(0, 6)
        self.grid.vert_wall(6, 0)
        self.grid.horz_wall(0, 12)
        self.grid.vert_wall(12, 0)
        # self.grid.horz_wall(0, 4)
        # self.grid.vert_wall(4, 0)

        # # gates
        self.grid.set(3, 6, None)
        self.grid.set(6, 3, None)
        self.grid.set(3, 12, None)
        self.grid.set(12, 3, None)

        self.grid.set(9, 6, None)
        self.grid.set(6, 9, None)
        self.grid.set(9, 12, None)
        self.grid.set(12, 9, None)

        self.grid.set(15, 12, None)
        self.grid.set(15, 6, None)

        self.grid.set(12, 15, None)
        self.grid.set(6, 15, None)


        # Randomize the player start position and orientation
        if self._agent_default_pos is not None:
            self.agent_pos = self._agent_default_pos
            self.grid.set(*self._agent_default_pos, None)
            self.agent_dir = 0 #self._rand_int(0, 4)  # assuming random start direction
        else:
            self.place_agent()

        if self._goal_default_pos is not None:
            goal = Goal()
            self.put_obj(goal, *self._goal_default_pos)
            goal.init_pos, goal.cur_pos = self._goal_default_pos

        self.mission = 'Reach the goal'

    def step(self, action):
        obs, reward, done, info = MiniGridEnv.step(self, action)
        return obs, reward, done, info


class NineRoomsDetEnv_v2(MiniGridEnv):
    """
    Classic 4 rooms gridworld environment.
    Can specify agent and goal position, if not it set at random.
    """

    def __init__(self, agent_pos=None, goal_pos=None):
        self._agent_default_pos = agent_pos
        self._goal_default_pos = goal_pos
        super().__init__(grid_size=19, max_steps=1000)    # always neet to set this up to more then the wraper one

    def _gen_grid(self, width, height):
        # Create the grid
        self.grid = Grid(width, height)

        # Generate the surrounding walls
        self.grid.horz_wall(0, 0)
        self.grid.horz_wall(0, height - 1)
        self.grid.vert_wall(0, 0)
        self.grid.vert_wall(width - 1, 0)

        # 9 rooms

        # #original
        # self.grid.horz_wall(0, 6)
        # self.grid.vert_wall(6, 0)
        # self.grid.horz_wall(0, 12)
        # self.grid.vert_wall(12, 0)
        #
        # # # gates
        # self.grid.set(12, 3, None)
        # self.grid.set(6, 3, None)
        # self.grid.set(3, 6, None)
        #
        # self.grid.set(9, 6, None)
        # self.grid.set(9, 14, None)
        # self.grid.set(6, 9, None)
        # self.grid.set(9, 12, None)
        # self.grid.set(12, 9, None)
        #
        # self.grid.set(15, 12, None)
        # self.grid.set(15, 6, None)
        #
        # self.grid.set(3, 12, None)
        # self.grid.set(12, 15, None)
        # self.grid.set(6, 15, None)

        #1 setup
        self.grid.horz_wall(0, 4)
        self.grid.vert_wall(6, 0)
        self.grid.horz_wall(0, 12)
        self.grid.vert_wall(12, 0)

        # # gates
        self.grid.set(3, 4, None)
        self.grid.set(8, 4, None)
        self.grid.set(3, 12, None)
        self.grid.set(12, 3, None)
        self.grid.set(6, 2, None)
        self.grid.set(15, 4, None)
        self.grid.set(15, 12, None)
        self.grid.set(12, 15, None)
        self.grid.set(6, 15, None)
        self.grid.set(6, 9, None)
        self.grid.set(9, 12, None)
        self.grid.set(12, 9, None)

        # # # 2 setup
        # self.grid.horz_wall(0, 4, length=7)
        # self.grid.horz_wall(0, 12, length=7)
        # self.grid.horz_wall(6, 6, length=7)
        # self.grid.horz_wall(6, 14, length=7)
        # self.grid.horz_wall(12, 4, length=7)
        # self.grid.horz_wall(12, 12, length=7)
        # self.grid.vert_wall(6, 0)
        # self.grid.vert_wall(12, 0)
        #
        # # # gates
        # self.grid.set(3, 4, None)
        # self.grid.set(3, 12, None)
        # self.grid.set(12, 3, None)
        # self.grid.set(6, 2, None)
        # self.grid.set(15, 4, None)
        # self.grid.set(9, 6, None)
        # self.grid.set(9, 14, None)
        # self.grid.set(6, 9, None)
        # self.grid.set(12, 9, None)
        # self.grid.set(12, 15, None)
        # self.grid.set(6, 15, None)

        # # 3 setup
        # self.grid.horz_wall(0, 3, length=3)
        # self.grid.vert_wall(3, 0, length=4)
        # self.grid.horz_wall(0, 6, length=3)
        # self.grid.vert_wall(3, 3, length=4)
        # self.grid.horz_wall(4, 3, length=3)
        # self.grid.vert_wall(6, 0, length=4)
        # self.grid.horz_wall(4, 6, length=3)
        # self.grid.vert_wall(6, 3, length=4)
        # self.grid.horz_wall(2, 9, length=2)
        # self.grid.vert_wall(3, 8, length=2)
        # self.grid.horz_wall(5, 9, length=2)
        # self.grid.horz_wall(0, 9)
        # self.grid.vert_wall(6, 0)
        #
        #
        # # # gates
        # self.grid.set(3, 1, None)
        # self.grid.set(3, 4, None)
        # self.grid.set(1, 3, None)
        # self.grid.set(4, 3, None)
        # self.grid.set(1, 9, None)
        # self.grid.set(4, 9, None)
        # self.grid.set(6, 7, None)
        # self.grid.set(4, 6, None)
        # self.grid.set(6, 4, None)
        # self.grid.set(6, 1, None)
        # self.grid.set(1, 6, None)
        # self.grid.set(12, 9, None)
        # self.grid.set(6, 12, None)


        # Randomize the player start position and orientation
        if self._agent_default_pos is not None:
            self.agent_pos = self._agent_default_pos
            self.grid.set(*self._agent_default_pos, None)
            self.agent_dir = 0 #self._rand_int(0, 4)  # assuming random start direction
        else:
            self.place_agent()

        if self._goal_default_pos is not None:
            goal = Goal()
            self.put_obj(goal, *self._goal_default_pos)
            goal.init_pos, goal.cur_pos = self._goal_default_pos

        self.mission = 'Reach the goal'

    def step(self, action):
        obs, reward, done, info = MiniGridEnv.step(self, action)
        return obs, reward, done, info

class NoRoomsDetEnv(NoRoomsDetEnv):
    def __init__(self, **kwargs):
        super().__init__(agent_pos=None, goal_pos=None)

class NoRoomsDetEnvGoalFixed1(NoRoomsDetEnv):
    def __init__(self, **kwargs):
        super().__init__(agent_pos=None, goal_pos=(17,17))

class SmallNoRoomsDetEnvGoalFixed1(SmallNoRoomsDetEnv):
    def __init__(self, **kwargs):
        super().__init__(agent_pos=(1, 1), goal_pos=(7,6))

class NoRoomsDetEnvGoalFixed2(NoRoomsDetEnv):
    def __init__(self, **kwargs):
        super().__init__(agent_pos=None, goal_pos=(1,17))

class NoRoomsDetEnvGoalFixed3(NoRoomsDetEnv):
    def __init__(self, **kwargs):
        super().__init__(agent_pos=None, goal_pos=(17,1))

class FourRoomsDetEnvGoalFixed0(FourRoomsDetEnv):
    def __init__(self, **kwargs):
        super().__init__(agent_pos=None, goal_pos=None)

class FourRoomsDetEnvGoalFixed1(FourRoomsDetEnv):
    def __init__(self, **kwargs):
        super().__init__(agent_pos=None, goal_pos=(17,17))

class NineRoomsDetEnvGoalFixed0(NineRoomsDetEnv):
    def __init__(self, goal_pos=None):
        super().__init__(agent_pos=None, goal_pos=None)

class NineRoomsDetEnvGoalFixed0v2(NineRoomsDetEnv_v2):
    def __init__(self, goal_pos=None):
        super().__init__(agent_pos=None, goal_pos=None)

class NineRoomsDetEnvGoalFixed1(NineRoomsDetEnv):
    def __init__(self, goal_pos=None):
        super().__init__(agent_pos=(1, 1), goal_pos=(5,17))

class FourRoomsDetEnvGoalFixed2(FourRoomsDetEnv):
    def __init__(self, **kwargs):
        super().__init__(agent_pos=None, goal_pos=(1,17))

class NineRoomsDetEnvGoalFixed2(NineRoomsDetEnv):
    def __init__(self, goal_pos=None):
        super().__init__(agent_pos=(1, 1), goal_pos=(11,17))

class FourRoomsDetEnvGoalFixed3(FourRoomsDetEnv):
    def __init__(self, **kwargs):
        super().__init__(agent_pos=None, goal_pos=(17,1))

class NineRoomsDetEnvGoalFixed3(NineRoomsDetEnv):
    def __init__(self, goal_pos=None):
        super().__init__(agent_pos=(1, 1), goal_pos=(17,17))
        
class NineRoomsDetEnvGoalFixed4(NineRoomsDetEnv):
    def __init__(self, goal_pos=None):
        super().__init__(agent_pos=(1, 1), goal_pos=(6,3))

register(
    id='MiniGrid-FourRooms-v0',
    entry_point='gym_minigrid.envs:FourRoomsEnv'
)

register(
    id='MiniGrid-FourRoomsDet-v0',
    entry_point='gym_minigrid.envs:FourRoomsDetEnvGoalFixed0'
)

register(
    id='MiniGrid-FourRoomsDet-v1',
    entry_point='gym_minigrid.envs:FourRoomsDetEnvGoalFixed1'
)

register(
    id='MiniGrid-NineRoomsDetv2-v0',
    entry_point='gym_minigrid.envs:NineRoomsDetEnvGoalFixed0v2'
)

register(
    id='MiniGrid-NineRoomsDet-v0',
    entry_point='gym_minigrid.envs:NineRoomsDetEnvGoalFixed0'
)

register(
    id='MiniGrid-NineRoomsDet-v1',
    entry_point='gym_minigrid.envs:NineRoomsDetEnvGoalFixed1'
)

register(
    id='MiniGrid-FourRoomsDet-v2',
    entry_point='gym_minigrid.envs:FourRoomsDetEnvGoalFixed2'
)

register(
    id='MiniGrid-NineRoomsDet-v2',
    entry_point='gym_minigrid.envs:NineRoomsDetEnvGoalFixed2'
)

register(
    id='MiniGrid-FourRoomsDet-v3',
    entry_point='gym_minigrid.envs:FourRoomsDetEnvGoalFixed3'
)

register(
    id='MiniGrid-NineRoomsDet-v3',
    entry_point='gym_minigrid.envs:NineRoomsDetEnvGoalFixed3'
)

register(
    id='MiniGrid-NineRoomsDet-v4',
    entry_point='gym_minigrid.envs:NineRoomsDetEnvGoalFixed4'
)

register(
    id='MiniGrid-NoRoomsDet-v0',
    entry_point='gym_minigrid.envs:NoRoomsDetEnv'
)

register(
    id='MiniGrid-NoRoomsDet-v1',
    entry_point='gym_minigrid.envs:NoRoomsDetEnvGoalFixed1'
)

register(
    id='MiniGrid-NoRoomsDet-v2',
    entry_point='gym_minigrid.envs:NoRoomsDetEnvGoalFixed2'
)

register(
    id='MiniGrid-NoRoomsDet-v3',
    entry_point='gym_minigrid.envs:NoRoomsDetEnvGoalFixed3'
)

register(
    id='MiniGrid-SmallNoRoomsDet-v1',
    entry_point='gym_minigrid.envs:SmallNoRoomsDetEnvGoalFixed1'
)