import random

import gym
from gym_minigrid.envs import Key, Ball, Box, Wall
from .verifier import *
from .levelgen import *


class Level_OpenRedDoor(RoomGridLevel):
    """
    Go to the red door
    (always unlocked, in the current room)
    Note: this level is intentionally meant for debugging and is
    intentionally kept very simple.
    """

    def __init__(self, seed=None):
        super().__init__(
            num_rows=1,
            num_cols=2,
            room_size=5,
            seed=seed
        )

    def gen_mission(self):
        obj, _ = self.add_door(0, 0, 0, 'red', locked=False)
        self.place_agent(0, 0)
        self.instrs = OpenInstr(ObjDesc('door', 'red'))

# class Level_BlockedUnlockPickup(RoomGridLevel):
#     """
#     Unlock a door blocked by a ball, then pick up a box
#     in another room
#     """
#
#     def __init__(self, seed=None):
#         room_size = 6
#         super().__init__(
#             num_rows=1,
#             num_cols=2,
#             room_size=room_size,
#             max_steps=16*room_size**2,
#             seed=seed
#         )
#
#     def gen_mission(self):
#         # Add a box to the room on the right
#         obj, _ = self.add_object(1, 0, kind="box")
#         # Make sure the two rooms are directly connected by a locked door
#         door, pos = self.add_door(0, 0, 0, locked=True)
#         # Block the door with a ball
#         color = self._rand_color()
#         self.grid.set(pos[0]-1, pos[1], Ball(color))
#         # Add a key to unlock the door
#         self.add_object(0, 0, 'key', door.color)
#
#         self.place_agent(0, 0)
#
#         self.instrs = PickupInstr(ObjDesc(obj.type))

class Level_TEST(RoomGridLevel):
    """
    Go to the red door
    (always unlocked, in the current room)
    Note: this level is intentionally meant for debugging and is
    intentionally kept very simple.
    """

    def __init__(self, seed=None):
        super().__init__(
            num_rows=1,
            num_cols=2,
            room_size=7,
            seed=seed
        )

    def gen_mission(self):
        # obj, _ = self.add_object(0, 0, kind="box")
        obj, _ = self.add_door(0, 0, 0, 'red', locked=True)
        # self.add_object(0, 0, 'key', 'red')
        self.grid.set(1,1, Ball('blue'))
        self.grid.set(2,4, Key('red'))
        self.grid.set(5,5, Box('green'))
        self.place_agent(0, 0)
        self.instrs = OpenInstr(ObjDesc('door', 'red'))



class Level_OpenDoor(RoomGridLevel):
    """
    Go to the door
    The door to open is given by its color or by its location.
    (always unlocked, in the current room)
    """

    def __init__(
        self,
        debug=False,
        select_by=None,

        seed=None
    ):
        self.select_by = select_by
        self.debug = debug
        super().__init__(seed=seed)

    def gen_mission(self):
        door_colors = self._rand_subset(COLOR_NAMES, 4)
        objs = []

        for i, color in enumerate(door_colors):
            obj, _ = self.add_door(1, 1, door_idx=i, color=color, locked=False)
            objs.append(obj)

        select_by = self.select_by
        if select_by is None:
            select_by = self._rand_elem(["color", "loc"])
        if select_by == "color":
            object = ObjDesc(objs[0].type, color=objs[0].color)
        elif select_by == "loc":
            object = ObjDesc(objs[0].type, loc=self._rand_elem(LOC_NAMES))

        self.place_agent(1, 1)
        self.instrs = OpenInstr(object, strict=self.debug)


class Level_OpenDoorDebug(Level_OpenDoor):
    """
    Same as OpenDoor but the level stops when any door is opened
    """

    def __init__(
        self,
        select_by=None,
        seed=None
    ):
        super().__init__(select_by=select_by, debug=True, seed=seed)


class Level_OpenDoorColor(Level_OpenDoor):
    """
    Go to the door
    The door is selected by color.
    (always unlocked, in the current room)
    """

    def __init__(self, seed=None):
        super().__init__(
            select_by="color",
            seed=seed
        )


#class Level_OpenDoorColorDebug(Level_OpenDoorColor, Level_OpenDoorDebug):
    """
    Same as OpenDoorColor but the level stops when any door is opened
    """
#    pass


class Level_OpenDoorLoc(Level_OpenDoor):
    """
    Go to the door
    The door is selected by location.
    (always unlocked, in the current room)
    """

    def __init__(self, seed=None):
        super().__init__(
            select_by="loc",
            seed=seed
        )


class Level_GoToDoor(RoomGridLevel):
    """
    Go to a door
    (of a given color, in the current room)
    No distractors, no language variation
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=7,
            seed=seed
        )

    def gen_mission(self):
        objs = []
        for _ in range(4):
            door, _ = self.add_door(1, 1)
            objs.append(door)
        self.place_agent(1, 1)

        obj = self._rand_elem(objs)
        self.instrs = GoToInstr(ObjDesc('door', obj.color))


class Level_GoToObjDoor(RoomGridLevel):
    """
    Go to an object or door
    (of a given type and color, in the current room)
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=8,
            seed=seed
        )

    def gen_mission(self):
        self.place_agent(1, 1)
        objs = self.add_distractors(1, 1, num_distractors=8, all_unique=False)

        for _ in range(4):
            door, _ = self.add_door(1, 1)
            objs.append(door)

        self.check_objs_reachable()

        obj = self._rand_elem(objs)
        self.instrs = GoToInstr(ObjDesc(obj.type, obj.color))


class Level_ActionObjDoor(RoomGridLevel):
    """
    [pick up an object] or
    [go to an object or door] or
    [open a door]
    (in the current room)
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=7,
            seed=seed
        )

    def gen_mission(self):
        objs, pos_all = self.add_distractors_withpos(1, 1, num_distractors=3)
        # objs, pos_all = self.add_distractors(1, 1, num_distractors=3)

        for _ in range(2):
            door, pos = self.add_door(1, 1, locked=False)
            objs.append(door)
            pos_all.append(pos)

        self.place_agent(1, 1)

        # obj = self._rand_elem(objs)
        obj, idx = self._rand_elem_withidx(objs)
        self.goal_pos = pos_all[idx]

        desc = ObjDesc(obj.type, obj.color)

        if obj.type == 'door':
            if self._rand_bool():
                self.instrs = GoToInstr(desc)
            else:
                self.instrs = OpenInstr(desc)
        else:
            if self._rand_bool():
                self.instrs = GoToInstr(desc)
            else:
                self.instrs = PickupInstr(desc)

################################################################################


class Level_TransDoor(RoomGridLevel):
    """
    [pick up an object] or
    [open a door]
    (in the current room)
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=7,
            seed=seed
        )

    def gen_mission(self):
        objs = []
        door_idx = random.randint(0,3)
        door, pos = self.add_door2(1, 1, door_idx=door_idx, color='yellow', locked=False)
        objs.append(door)


        self.place_agent(1, 1)
        desc_all = []
        instrs_all = []
        instrs_all_addition = []
        # print('task-generate begin')

        obj = self._rand_elem(objs)

        desc = ObjDesc(obj.type, obj.color)

        if obj.type == 'door':
            self.instrs = OpenInstr(desc)
        else:
            self.instrs = PickupInstr(desc)

        self.instrs_all = instrs_all
        self.instrs_all_addition = instrs_all_addition


class Level_TransObj(RoomGridLevel):
    """
    [pick up an object] or
    [open a door]
    (in the current room)
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=7,
            seed=seed
        )

    def gen_mission(self):
        objs = []
        # door_idx = random.randint(0,3)
        # door, pos = self.add_door2(1, 1, door_idx=door_idx, color='yellow', locked=False)
        obj, pos = self.add_object(1, 1, kind='ball', color='yellow')
        self.goal_pos_obj = pos
        objs.append(obj)


        self.place_agent(1, 1)
        desc_all = []
        instrs_all = []
        instrs_all_addition = []
        # print('task-generate begin')

        obj = self._rand_elem(objs)

        desc = ObjDesc(obj.type, obj.color)

        if obj.type == 'door':
            self.instrs = OpenInstr(desc)
        else:
            self.instrs = PickupInstr(desc)

        self.instrs_all = instrs_all
        self.instrs_all_addition = instrs_all_addition

##########################################

class Level_MyActionObjDoor(RoomGridLevel):
    """
    [pick up an object] or
    [open a door]
    (in the current room)
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=7,
            seed=seed
        )

    def gen_mission(self):
        objs_1 = self.add_distractors(0, 1, num_distractors=1)
        objs_2 = self.add_distractors(1, 0, num_distractors=1)
        objs_3 = self.add_distractors(2, 1, num_distractors=1)
        objs_4 = self.add_distractors(1, 2, num_distractors=1)
        objs_addition = [objs_1,objs_2,objs_3, objs_4]

        objs = self.add_distractors(1, 1, num_distractors=2)
        for _ in range(2):
            door, _ = self.add_door(1, 1, locked=False)
            objs.append(door)

        for _ in range(2):
            door, _ = self.add_door(1, 1, locked=True)
            key, _ = self.add_object(1, 1, 'key', door.color)
            objs.append(key)
            objs.append(door)



        self.place_agent(1, 1)
        desc_all = []
        instrs_all = []
        instrs_all_addition = []
        # print('task-generate begin')
        for obj in objs:
            # print(obj.type, obj.color)
            desc = ObjDesc(obj.type, obj.color)
            desc_all.append(desc)
            if obj.type == 'door':
                instrs_all.append(OpenInstr(desc))
                # print(OpenInstr(desc).surface(self))
            else:
                instrs_all.append(PickupInstr(desc))
                # print(PickupInstr(desc).surface(self))

        for item in objs_addition:
            # print(obj.type, obj.color)
            desc_1 = ObjDesc(item[0].type, item[0].color)
            if item[0].type == 'door':
                instrs_all_addition.append(OpenInstr(desc_1))
                # print(OpenInstr(desc).surface(self))
            else:
                instrs_all_addition.append(PickupInstr(desc_1))
                # print(PickupInstr(desc).surface(self))



        obj = self._rand_elem(objs)

        desc = ObjDesc(obj.type, obj.color)

        if obj.type == 'door':
            self.instrs = OpenInstr(desc)
        else:
            self.instrs = PickupInstr(desc)
        self.instrs_all = instrs_all
        self.instrs_all_addition = instrs_all_addition
        # self.instrs = instrs_all


class Level_InsVrf(RoomGridLevel):
    """
    [pick up an object] or
    [open a door]
    (in the current room)
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=7,
            seed=seed
        )

    def gen_mission(self):
        # objs, pos_all = self.add_distractors_withpos(1, 1, num_distractors=2)
        # objs, pos_all = self.add_distractors(1, 1, num_distractors=3)
        objs_, pos_ = self.add_object(1, 1)

        objs = []
        pos_all = []
        self.room_objs_recorder = []
        # obj_color = random.choice(['red', 'yellow', 'green', 'blue', ])

        if random.random() > 0.5:
            obj, pos = self.add_object(1, 1, kind='ball', color='red')
            objs.append(obj)
            pos_all.append(pos)
            self.room_objs_recorder.append([obj.color, obj.type])



        # obj_color = random.choice(['red', 'yellow', 'green', 'blue',  ])

        if random.random() > 0.5:
            obj, pos = self.add_object(1, 1, kind='key', color='blue')
            objs.append(obj)
            pos_all.append(pos)
            self.room_objs_recorder.append([obj.color, obj.type])



        door_num = 1


        for _ in range(door_num):
            door_color = random.choice(['yellow',])
            if random.random() > 0.:
                door, pos = self.add_door(1, 1, color = door_color, locked=False)
                objs.append(door)
                pos_all.append(pos)
                self.room_objs_recorder.append([door.color, door.type])

            else:
                door, pos = self.add_door(1, 1, color = door_color, locked=True)
                key, pos_key = self.add_object(1, 1, 'key', door.color)
                objs.append(door)
                pos_all.append(pos)
                objs.append(key)
                pos_all.append(pos_key)
                self.room_objs_recorder.append([door.color, door.type])
                self.room_objs_recorder.append([key.color, key.type])

        self.place_agent(1, 1)

        # obj = self._rand_elem(objs)
        obj, idx = self._rand_elem_withidx(objs)
        self.goal_pos = pos_all[idx]

        desc = ObjDesc(obj.type, obj.color)

        if obj.type == 'door':
            if self._rand_bool():
                self.instrs = GoToInstr(desc)
            else:
                self.instrs = OpenInstr(desc)
        else:
            if self._rand_bool():
                self.instrs = GoToInstr(desc)
            else:
                self.instrs = PickupInstr(desc)





class Level_ActionOneObj(RoomGridLevel):
    """
    [pick up an object] or
    [open a door]
    (in the current room)
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=7,
            seed=seed
        )

    def gen_mission(self):

        # objs = self.add_distractors(1, 1, num_distractors=1)
        obj_list = self.add_object_norand(1, 1, kind='ball',color='red')

        objs = []
        objs.append(obj_list[0])
        # print(obj_list[1])
        # pos are x,y in [7,11]



        # for _ in range(2):
        #     door, _ = self.add_door(1, 1, locked=False)
        #     objs.append(door)

        # for _ in range(2):
        #     door, _ = self.add_door(1, 1, locked=True)
        #     key, _ = self.add_object(1, 1, 'key', door.color)
        #     objs.append(key)
        #     objs.append(door)



        self.place_agent(1, 1)
        desc_all = []
        instrs_all = []
        instrs_all_addition = []
        # print('task-generate begin')
        for obj in objs:
            # print(obj.type, obj.color)
            desc = ObjDesc(obj.type, obj.color)
            desc_all.append(desc)
            if obj.type == 'door':
                instrs_all.append(OpenInstr(desc))
                # print(OpenInstr(desc).surface(self))
            else:
                instrs_all.append(PickupInstr(desc))
                # print(PickupInstr(desc).surface(self))

        # for item in objs_addition:
        #     # print(obj.type, obj.color)
        #     desc_1 = ObjDesc(item[0].type, item[0].color)
        #     if item[0].type == 'door':
        #         instrs_all_addition.append(OpenInstr(desc_1))
        #         # print(OpenInstr(desc).surface(self))
        #     else:
        #         instrs_all_addition.append(PickupInstr(desc_1))
        #         # print(PickupInstr(desc).surface(self))



        obj = self._rand_elem(objs)

        desc = ObjDesc(obj.type, obj.color)

        if obj.type == 'door':
            self.instrs = OpenInstr(desc)
        else:
            self.instrs = PickupInstr(desc)


class Level_ActionMltObj(RoomGridLevel):
    """
    [pick up an object] or
    [open a door]
    (in the current room)
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=7,
            seed=seed
        )

    def gen_mission(self):

        # objs = self.add_distractors(1, 1, num_distractors=1)
        obj_list = self.add_object_norand(1, 1, kind='ball',color='red')
        self.goal_pos = obj_list[1]
        objs = []
        objs.append(obj_list[0])

        self.place_agent(1, 1)
        desc_all = []
        instrs_all = []
        instrs_all_addition = []
        # print(self.agent_pos)
        self.agent_pos = [9,9]
        number = random.randint(1,3)
        obj_add = self.add_distractors(1, 1, num_distractors=number)

        obj = self._rand_elem(objs)

        desc = ObjDesc(obj.type, obj.color)

        if obj.type == 'door':
            self.instrs = OpenInstr(desc)
        else:
            self.instrs = PickupInstr(desc)

class Level_ActionStrObj(RoomGridLevel):
    """
    [pick up an object] or
    [open a door]
    (in the current room)
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=7,
            seed=seed
        )

    def gen_mission(self):

        # objs = self.add_distractors(1, 1, num_distractors=1)
        obj_list = self.add_object_norand(1, 1, kind='ball',color='red')

        objs = []
        objs.append(obj_list[0])

        self.place_agent(1, 1)
        self.agent_pos = [11,11]
        wall_length = random.randint(1,3)
        for j in range(wall_length):
            self.grid.set(7+j,9, Wall())

        self.grid.set(7 + j, 9 + 1, Wall())
        self.grid.set(7 + j, 9 + 2, Wall())

        desc_all = []
        instrs_all = []
        instrs_all_addition = []


        obj = self._rand_elem(objs)

        desc = ObjDesc(obj.type, obj.color)

        if obj.type == 'door':
            self.instrs = OpenInstr(desc)
        else:
            self.instrs = PickupInstr(desc)


class Level_aba_nopos(RoomGridLevel):
    """
    [pick up an object] or
    [open a door]
    (in the current room)
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=7,
            seed=seed
        )

    def gen_mission(self):

        # objs = self.add_distractors(1, 1, num_distractors=1)
        obj_list = self.add_object_norand(1, 1, kind='ball', color='red', put_x=8, put_y=8)

        if random.random() < 0.5:
            if random.random() < 1:
                door, pos = self.add_door_norand(1, 1, locked=False)
            elif random.random() < 0.66:
                door, pos = self.add_door(1, 1, door_idx=0, locked=False)
            else:
                door, pos = self.add_door(1, 1, door_idx=3, locked=False)

            self.place_agent(1, 1)
            self.agent_pos = [11,11]

            wall_length = random.randint(1, 3)
            # wall_length = 2
            for j in range(wall_length):
                self.grid.set(7 + j, 9, Wall())

            self.grid.set(7 + j, 9 + 1, Wall())
            self.grid.set(7 + j, 9 + 2, Wall())

        else:

            if random.random() < 1:
                door, pos = self.add_door_norand(1, 1, locked=False, pos=[8, 6])
            elif random.random() < 0.66:
                door, pos = self.add_door(1, 1, door_idx=1, locked=False)
            else:
                door, pos = self.add_door(1, 1, door_idx=2, locked=False)

            self.place_agent(1, 1)
            self.agent_pos = [11,11]

            wall_length = random.randint(1, 3)
            # print('build wall')
            # print(wall_length)
            for j in range(wall_length):
                self.grid.set(11 - j, 9, Wall())

            self.grid.set(11 - j, 9 - 1, Wall())
            self.grid.set(11 - j, 9 - 2, Wall())

        objs = []
        objs.append(obj_list[0])
        objs.append(door)

        desc_all = []
        instrs_all = []
        instrs_all_addition = []

        obj = self._rand_elem(objs)

        number = random.randint(1, 2)
        obj_add = self.add_distractors(1, 1, num_distractors=number)

        desc = ObjDesc(obj.type, obj.color)

        if obj.type == 'door':
            self.instrs = OpenInstr(desc)
        else:
            self.instrs = PickupInstr(desc)

class Level_aba_nopos2(RoomGridLevel):
    """
    [pick up an object] or
    [open a door]
    (in the current room)
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=7,
            seed=seed
        )

    def gen_mission(self):

        # objs = self.add_distractors(1, 1, num_distractors=1)
        obj_list = self.add_object_norand(1, 1, kind='ball', color='red', put_x=10, put_y=10)

        if random.random() < 0.5:
            if random.random() < 1:
                door, pos = self.add_door_norand(1, 1, locked=False)
            elif random.random() < 0.66:
                door, pos = self.add_door(1, 1, door_idx=0, locked=False)
            else:
                door, pos = self.add_door(1, 1, door_idx=3, locked=False)

            self.place_agent(1, 1)
            self.agent_pos = [7,7]

            wall_length = random.randint(1, 3)
            # wall_length = 2
            for j in range(wall_length):
                self.grid.set(7 + j, 9, Wall())

            self.grid.set(7 + j, 9 + 1, Wall())
            self.grid.set(7 + j, 9 + 2, Wall())

        else:

            if random.random() < 1:
                door, pos = self.add_door_norand(1, 1, locked=False, pos=[8, 6])
            elif random.random() < 0.66:
                door, pos = self.add_door(1, 1, door_idx=1, locked=False)
            else:
                door, pos = self.add_door(1, 1, door_idx=2, locked=False)

            self.place_agent(1, 1)
            self.agent_pos = [11,11]

            wall_length = random.randint(1, 3)
            # print('build wall')
            # print(wall_length)
            for j in range(wall_length):
                self.grid.set(11 - j, 9, Wall())

            self.grid.set(11 - j, 9 - 1, Wall())
            self.grid.set(11 - j, 9 - 2, Wall())

        objs = []
        objs.append(obj_list[0])
        objs.append(door)

        desc_all = []
        instrs_all = []
        instrs_all_addition = []

        obj = self._rand_elem(objs)

        number = random.randint(1, 2)
        obj_add = self.add_distractors(1, 1, num_distractors=number)

        desc = ObjDesc(obj.type, obj.color)

        if obj.type == 'door':
            self.instrs = OpenInstr(desc)
        else:
            self.instrs = PickupInstr(desc)




class Level_aba_nointer(RoomGridLevel):
    """
    [pick up an object] or
    [open a door]
    (in the current room)
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=7,
            seed=seed
        )

    def gen_mission(self):

        # objs = self.add_distractors(1, 1, num_distractors=1)
        obj_list = self.add_object_norand(1, 1, kind='ball',color='red',put_x=8,put_y=8)
        if random.random()<0.5:
            if random.random() < 0.33:
                door, pos = self.add_door_norand(1, 1, locked=False)
            elif random.random() < 0.66:
                door, pos = self.add_door(1, 1, door_idx=0, locked=False)
            else:
                door, pos = self.add_door(1, 1, door_idx=3, locked=False)


            self.place_agent(1, 1)
            self.agent_pos = [random.randint(10, 11), random.randint(7, 11)]

            wall_length = random.randint(1, 3)

            for j in range(wall_length):
                self.grid.set(7 + j, 9, Wall())

            self.grid.set(7 + j, 9 + 1, Wall())
            self.grid.set(7 + j, 9 + 2, Wall())

        else:

            if random.random() < 0.33:
                door, pos = self.add_door_norand(1, 1, locked=False, pos=[8,6])
            elif random.random() < 0.66:
                door, pos = self.add_door(1, 1, door_idx=1, locked=False)
            else:
                door, pos = self.add_door(1, 1, door_idx=2, locked=False)


            self.place_agent(1, 1)
            self.agent_pos = [random.randint(7, 11), random.randint(10, 11)]

            wall_length = random.randint(1, 3)
            # print('build wall')
            # print(wall_length)
            for j in range(wall_length):
                self.grid.set(11 - j, 9, Wall())

            self.grid.set(11 - j, 9 - 1, Wall())
            self.grid.set(11 - j, 9 - 2, Wall())


        objs = []
        objs.append(obj_list[0])
        objs.append(door)




        desc_all = []
        instrs_all = []
        instrs_all_addition = []


        obj = self._rand_elem(objs)

        desc = ObjDesc(obj.type, obj.color)

        if obj.type == 'door':
            self.instrs = OpenInstr(desc)
        else:
            self.instrs = PickupInstr(desc)





class Level_aba_nostr(RoomGridLevel):
    """
    [pick up an object] or
    [open a door]
    (in the current room)
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=7,
            seed=seed
        )

    def gen_mission(self):

        # objs = self.add_distractors(1, 1, num_distractors=1)
        obj_list = self.add_object_norand(1, 1, kind='ball', color='red', put_x=8, put_y=8)

        if random.random() < 1:
            if random.random() < 0.33:
                door, pos = self.add_door_norand(1, 1, locked=False)
            elif random.random() < 0.66:
                door, pos = self.add_door(1, 1, door_idx=0, locked=False)
            else:
                door, pos = self.add_door(1, 1, door_idx=3, locked=False)

            self.place_agent(1, 1)
            self.agent_pos = [random.randint(10, 11), random.randint(7, 11)]

            # wall_length = random.randint(1, 3)
            wall_length = 2
            for j in range(wall_length):
                self.grid.set(7 + j, 9, Wall())

            self.grid.set(7 + j, 9 + 1, Wall())
            self.grid.set(7 + j, 9 + 2, Wall())

        else:
            pass
            if random.random() < 0.33:
                door, pos = self.add_door_norand(1, 1, locked=False, pos=[8, 6])
            elif random.random() < 0.66:
                door, pos = self.add_door(1, 1, door_idx=1, locked=False)
            else:
                door, pos = self.add_door(1, 1, door_idx=2, locked=False)

            self.place_agent(1, 1)
            self.agent_pos = [random.randint(7, 11), random.randint(10, 11)]

            wall_length = random.randint(1, 3)
            # print('build wall')
            # print(wall_length)
            for j in range(wall_length):
                self.grid.set(11 - j, 9, Wall())

            self.grid.set(11 - j, 9 - 1, Wall())
            self.grid.set(11 - j, 9 - 2, Wall())

        objs = []
        objs.append(obj_list[0])
        objs.append(door)

        desc_all = []
        instrs_all = []
        instrs_all_addition = []

        obj = self._rand_elem(objs)


        number = random.randint(1,2)
        obj_add = self.add_distractors(1, 1, num_distractors=number)


        desc = ObjDesc(obj.type, obj.color)

        if obj.type == 'door':
            self.instrs = OpenInstr(desc)
        else:
            self.instrs = PickupInstr(desc)


class Level_aba_nostr2(RoomGridLevel):
    """
    [pick up an object] or
    [open a door]
    (in the current room)
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=7,
            seed=seed
        )

    def gen_mission(self):

        # objs = self.add_distractors(1, 1, num_distractors=1)
        obj_list = self.add_object_norand(1, 1, kind='ball', color='red', put_x=10, put_y=10)

        if random.random() < 1:
            if random.random() < 0.33:
                door, pos = self.add_door_norand(1, 1, locked=False, pos=(10,12))
            elif random.random() < 0.66:
                door, pos = self.add_door(1, 1, door_idx=0, locked=False)
            else:
                door, pos = self.add_door(1, 1, door_idx=3, locked=False)

            self.place_agent(1, 1)
            self.agent_pos = [random.randint(7, 11), random.randint(7, 8)]

            # wall_length = random.randint(1, 3)
            wall_length = 2
            for j in range(wall_length):
                self.grid.set(7 + j, 9, Wall())

            self.grid.set(7 + j, 9 + 1, Wall())
            self.grid.set(7 + j, 9 + 2, Wall())

        else:
            pass
            if random.random() < 0.33:
                door, pos = self.add_door_norand(1, 1, locked=False, pos=[10,12])
            elif random.random() < 0.66:
                door, pos = self.add_door(1, 1, door_idx=1, locked=False)
            else:
                door, pos = self.add_door(1, 1, door_idx=2, locked=False)

            self.place_agent(1, 1)
            self.agent_pos = [random.randint(7, 8), random.randint(7,11)]

            wall_length = random.randint(1, 3)
            # print('build wall')
            # print(wall_length)
            for j in range(wall_length):
                self.grid.set(11 - j, 9, Wall())

            self.grid.set(11 - j, 9 - 1, Wall())
            self.grid.set(11 - j, 9 - 2, Wall())

        objs = []
        objs.append(obj_list[0])
        objs.append(door)

        desc_all = []
        instrs_all = []
        instrs_all_addition = []

        obj = self._rand_elem(objs)


        number = random.randint(1,2)
        obj_add = self.add_distractors(1, 1, num_distractors=number)


        desc = ObjDesc(obj.type, obj.color)

        if obj.type == 'door':
            self.instrs = OpenInstr(desc)
        else:
            self.instrs = PickupInstr(desc)





class Level_aba_noall(RoomGridLevel):
    """
    [pick up an object] or
    [open a door]
    (in the current room)
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=7,
            seed=seed
        )

    def gen_mission(self):

        # objs = self.add_distractors(1, 1, num_distractors=1)
        obj_list = self.add_object_norand(1, 1, kind='ball', color='red', put_x=8, put_y=8)

        if random.random() < 1:
            if random.random() < 1:
                door, pos = self.add_door_norand(1, 1, locked=False)
            elif random.random() < 0.66:
                door, pos = self.add_door(1, 1, door_idx=0, locked=False)
            else:
                door, pos = self.add_door(1, 1, door_idx=3, locked=False)

            self.place_agent(1, 1)
            self.agent_pos = [11,11]

            # wall_length = random.randint(1, 3)
            wall_length = 2
            for j in range(wall_length):
                self.grid.set(7 + j, 9, Wall())

            self.grid.set(7 + j, 9 + 1, Wall())
            self.grid.set(7 + j, 9 + 2, Wall())

        else:
            pass
            if random.random() < 0.33:
                door, pos = self.add_door_norand(1, 1, locked=False, pos=[8, 6])
            elif random.random() < 0.66:
                door, pos = self.add_door(1, 1, door_idx=1, locked=False)
            else:
                door, pos = self.add_door(1, 1, door_idx=2, locked=False)

            self.place_agent(1, 1)
            self.agent_pos = [random.randint(7, 11), random.randint(10, 11)]

            wall_length = random.randint(1, 3)
            # print('build wall')
            # print(wall_length)
            for j in range(wall_length):
                self.grid.set(11 - j, 9, Wall())

            self.grid.set(11 - j, 9 - 1, Wall())
            self.grid.set(11 - j, 9 - 2, Wall())

        objs = []
        objs.append(obj_list[0])
        objs.append(door)

        desc_all = []
        instrs_all = []
        instrs_all_addition = []

        obj = self._rand_elem(objs)

        # number = random.randint(1, 2)
        # obj_add = self.add_distractors(1, 1, num_distractors=number)

        desc = ObjDesc(obj.type, obj.color)

        if obj.type == 'door':
            self.instrs = OpenInstr(desc)
        else:
            self.instrs = PickupInstr(desc)



class Level_UnlockPickup2FixRoom(RoomGridLevel):
    """
    Unlock a door, then pick up a box in another room
    """

    def __init__(self, distractors=False, seed=None):
        self.distractors = distractors

        room_size = 7
        super().__init__(
            num_rows=1,
            num_cols=3,
            room_size=room_size,
            max_steps=8*room_size**2,
            seed=seed
        )

    def gen_mission(self):
        # Add a random object to the room on the right
        obj, pos_ball = self.add_object(1, 0, color='yellow', kind="ball")
        # Make sure the two rooms are directly connected by a locked door
        door, pos_door = self.add_door(0, 0, 0, color= 'red', locked=True)
        # Add a key to unlock the door
        _, pos_key = self.add_object(0, 0, 'key', door.color)

        # door, _ = self.add_door(1, 0, 0, color='blue', locked=True)
        # Add a key to unlock the door
        # self.add_object(1, 0, 'key', door.color)

        if self.distractors:
            self.add_distractors(num_distractors=4)

        self.place_agent(0, 0)

        self.instrs = PickupInstr(ObjDesc(obj.type, obj.color))
        self.goal_pos = []
        self.goal_pos.append(pos_key)
        self.goal_pos.append(pos_door)
        self.goal_pos.append(pos_ball)



class Level_UnlockPickup3FixRoom(RoomGridLevel):
    """
    Unlock a door, then pick up a box in another room
    """

    def __init__(self, distractors=False, seed=None):
        self.distractors = distractors

        room_size = 7
        super().__init__(
            num_rows=1,
            num_cols=3,
            room_size=room_size,
            max_steps=8*room_size**2,
            seed=seed
        )

    def gen_mission(self):
        # Add a random object to the room on the right
        obj, pos_ball = self.add_object(2, 0, color='yellow', kind="ball")
        # Make sure the two rooms are directly connected by a locked door
        door, pos_door_1 = self.add_door(0, 0, 0, color= 'red', locked=True)
        # Add a key to unlock the door
        _, pos_key_1 = self.add_object(0, 0, 'key', door.color)

        door, pos_door_2 = self.add_door(1, 0, 0, color='blue', locked=True)
        # Add a key to unlock the door
        _, pos_key_2 = self.add_object(1, 0, 'key', door.color)

        if self.distractors:
            self.add_distractors(num_distractors=4)

        self.place_agent(0, 0)

        self.instrs = PickupInstr(ObjDesc(obj.type, obj.color))

        self.goal_pos = []
        self.goal_pos.append(pos_key_1)
        self.goal_pos.append(pos_door_1)
        self.goal_pos.append(pos_key_2)
        self.goal_pos.append(pos_door_2)
        self.goal_pos.append(pos_ball)




class Level_UnlockPickup2FixRoomSTR(RoomGridLevel):
    """
    Unlock a door, then pick up a box in another room
    """
    def __init__(self, distractors=False, seed=None):
        self.distractors = distractors

        room_size = 7
        super().__init__(
            num_rows=1,
            num_cols=3,
            room_size=room_size,
            max_steps=8*room_size**2,
            seed=seed
        )

    def gen_mission(self):
        # Add a random object to the room on the right
        obj, pos_ball = self.add_object(1, 0, color='red', kind="ball")
        # Make sure the two rooms are directly connected by a locked door
        door, pos_door = self.add_door(0, 0, 0, color= 'red', locked=True)
        # Add a key to unlock the door
        _, pos_key = self.add_object(0, 0, 'key', door.color)

        # door, _ = self.add_door(1, 0, 0, color='blue', locked=True)
        # Add a key to unlock the door
        # self.add_object(1, 0, 'key', door.color)


        wall_length = random.randint(1, 2)


        wall_x_room1 = 3
        wall_x_room2 = 9
        wall_y_room1 = 3


        wall_x1 = random.randint(wall_x_room1-1,wall_x_room1+1)

        # print('build wall')
        # print(wall_length)
        if pos_ball[0] != wall_x_room1 and pos_key[0] != wall_x_room1:
            if random.random() < 0.5:
                for j in range(wall_length):
                    self.grid.set(wall_x1, 5 - j, Wall())
            else:
                for j in range(wall_length):
                    self.grid.set(wall_x1, 1 + j, Wall())

        elif pos_ball[1] != wall_y_room1 and pos_key[1] != wall_y_room1:
            for j in range(wall_length):
                self.grid.set(1+j ,wall_y_room1, Wall())










        if self.distractors:
            self.add_distractors(num_distractors=4)

        self.place_agent(0, 0)

        self.instrs = PickupInstr(ObjDesc(obj.type, obj.color))
        self.goal_pos = []
        self.goal_pos.append(pos_key)
        self.goal_pos.append(pos_door)
        self.goal_pos.append(pos_ball)

class Level_UnlockPickup2FixRoomSTR2(RoomGridLevel):
    """
    Unlock a door, then pick up a box in another room
    """
    def __init__(self, distractors=False, seed=None):
        self.distractors = distractors

        room_size = 7
        super().__init__(
            num_rows=1,
            num_cols=3,
            room_size=room_size,
            max_steps=8*room_size**2,
            seed=seed
        )

    def gen_mission(self):
        # Add a random object to the room on the right
        obj, pos_ball = self.add_object(1, 0, color='red', kind="ball")
        # Make sure the two rooms are directly connected by a locked door
        door, pos_door = self.add_door(0, 0, 0, color= 'red', locked=True)
        # Add a key to unlock the door
        _, pos_key = self.add_object(0, 0, 'key', door.color)

        # door, _ = self.add_door(1, 0, 0, color='blue', locked=True)
        # Add a key to unlock the door
        # self.add_object(1, 0, 'key', door.color)


        wall_length = random.randint(1, 2)


        wall_x_room1 = 3
        wall_x_room2 = 9
        wall_y_room1 = 3


        wall_x1 = random.randint(wall_x_room1-1,wall_x_room1)

        # print('build wall')
        # print(wall_length)
        if pos_ball[0] != wall_x_room1 and pos_key[0] != wall_x_room1:
            if random.random() < 0.5:
                for j in range(wall_length):
                    self.grid.set(wall_x1, 5 - j, Wall())
            else:
                for j in range(wall_length):
                    self.grid.set(wall_x1, 1 + j, Wall())

        elif pos_ball[1] != wall_y_room1 and pos_key[1] != wall_y_room1:
            for j in range(wall_length):
                self.grid.set(1+j ,wall_y_room1, Wall())










        if self.distractors:
            self.add_distractors(num_distractors=4)

        self.place_agent(0, 0)

        self.instrs = PickupInstr(ObjDesc(obj.type, obj.color))
        self.goal_pos = []
        self.goal_pos.append(pos_key)
        self.goal_pos.append(pos_door)
        self.goal_pos.append(pos_ball)





class Level_UnlockPickup2FixRoomSTR_Round(RoomGridLevel):
    """
    Unlock a door, then pick up a box in another room
    """
    def __init__(self, distractors=False, seed=None):
        self.distractors = distractors

        room_size = 7
        super().__init__(
            num_rows=3,
            num_cols=3,
            room_size=room_size,
            max_steps=8*room_size**2,
            seed=seed
        )

    def gen_mission(self):

        direction = random.randint(0,3)
        # 0 right; 1 down,  2 left;  3 up

        # Add a random object to the room on the right

        # room codx   0 0;  1 0;  2 0
        #             0 1;  1 1;  2 1
        #             0 2;  1 2;  2 2



        des_pos_list = [[7,7],[11,11],[11,7],[7,11],
                        [8,8],[10,10],[10,8],[8,10],
                        [7, 7], [10, 10], [10, 7], [7, 10],
                        [8, 8], [11, 11], [11, 8], [8, 11],
                        ]
        pos_select = random.sample(des_pos_list,2)
        door, pos_door = self.add_door(1, 1, door_idx=direction, color='red', locked=True)
        wall_x_room = 9
        wall_y_room = 9
        self.place_agent(1, 1)
        self.agent_pos = pos_select[1]

        # wall_length = random.randint(1, 2)
        wall_length = 2


        if direction == 0: # right
            obj, pos_ball = self.add_object_norand(2, 1, color='red', kind="ball")
            _, pos_key = self.add_object_norand(1, 1, 'key', door.color, put_x=pos_select[0][0], put_y=pos_select[0][1])
            for j in range(wall_length):
                self.grid.set(7+j, wall_y_room, Wall())




        elif direction == 1:  # down
            obj, pos_ball = self.add_object_norand(1, 2, color='red', kind="ball")
            _, pos_key = self.add_object_norand(1, 1, 'key', door.color, put_x=pos_select[0][0],
                                                put_y=pos_select[0][1])
            for j in range(wall_length):
                self.grid.set(wall_x_room, 7+j, Wall())



        elif direction == 2:  # left
            obj, pos_ball = self.add_object_norand(0, 1, color='red', kind="ball")
            _, pos_key = self.add_object_norand(1, 1, 'key', door.color, put_x=pos_select[0][0],
                                                put_y=pos_select[0][1])
            for j in range(wall_length):
                self.grid.set(11 - j, wall_y_room, Wall())



        else:  # up
            obj, pos_ball = self.add_object_norand(1, 0, color='red', kind="ball")
            _, pos_key = self.add_object_norand(1, 1, 'key', door.color, put_x=pos_select[0][0],
                                                put_y=pos_select[0][1])
            for j in range(wall_length):
                self.grid.set(wall_x_room, 11-j, Wall())





        # print('build wall')
        # print(wall_length)
        # if pos_ball[0] != wall_x_room1 and pos_key[0] != wall_x_room1:
        #     if random.random() < 0.5:
        #         for j in range(wall_length):
        #             self.grid.set(wall_x1, 5 - j, Wall())
        #     else:
        #         for j in range(wall_length):
        #             self.grid.set(wall_x1, 1 + j, Wall())
        #
        # elif pos_ball[1] != wall_y_room1 and pos_key[1] != wall_y_room1:
        #     for j in range(wall_length):
        #         self.grid.set(1+j ,wall_y_room1, Wall())










        if self.distractors:
            self.add_distractors(num_distractors=4)



        self.instrs = PickupInstr(ObjDesc(obj.type, obj.color))
        self.goal_pos = []
        self.goal_pos.append(pos_key)
        self.goal_pos.append(pos_door)
        self.goal_pos.append(pos_ball)


class Level_UnlockPickup3FixRoomSTR(RoomGridLevel):
    """
    Unlock a door, then pick up a box in another room
    """

    def __init__(self, distractors=False, seed=None):
        self.distractors = distractors

        room_size = 7
        super().__init__(
            num_rows=1,
            num_cols=3,
            room_size=room_size,
            max_steps=8*room_size**2,
            seed=seed
        )

    def gen_mission(self):
        # Add a random object to the room on the right
        obj, pos_ball = self.add_object(2, 0, color='yellow', kind="ball")
        # Make sure the two rooms are directly connected by a locked door
        door, pos_door_1 = self.add_door(0, 0, 0, color= 'red', locked=True)
        # Add a key to unlock the door
        _, pos_key_1 = self.add_object(0, 0, 'key', door.color)

        door, pos_door_2 = self.add_door(1, 0, 0, color='blue', locked=True)
        # Add a key to unlock the door
        _, pos_key_2 = self.add_object(1, 0, 'key', door.color)

        if self.distractors:
            self.add_distractors(num_distractors=4)

        self.place_agent(0, 0)

        self.instrs = PickupInstr(ObjDesc(obj.type, obj.color))

        self.goal_pos = []
        self.goal_pos.append(pos_key_1)
        self.goal_pos.append(pos_door_1)
        self.goal_pos.append(pos_key_2)
        self.goal_pos.append(pos_door_2)
        self.goal_pos.append(pos_ball)






class Level_UnlockPickup5FixRoom(RoomGridLevel):
    """
    Unlock a door, then pick up a box in another room
    """

    def __init__(self, distractors=False, seed=None):
        self.distractors = distractors

        room_size = 7
        super().__init__(
            num_rows=3,
            num_cols=3,
            room_size=room_size,
            max_steps=8*room_size**2,
            seed=seed
        )

    def gen_mission(self):
        # Add a random object to the room on the right
        obj, _ = self.add_object(2, 2, color= 'yellow', kind="ball")
        # Make sure the two rooms are directly connected by a locked door


        door, _ = self.add_door(0, 0, 0, color= 'red', locked=True)
        # Add a key to unlock the door
        self.add_object(0, 0, 'key', door.color)


        door, _ = self.add_door(1, 0, 0, color='yellow', locked=True)
        # Add a key to unlock the door
        self.add_object(1, 0, 'key', door.color)


        door, _ = self.add_door(2, 0, 1, color='grey', locked=True)
        # Add a key to unlock the door
        self.add_object(2, 0, 'key', door.color)


        door, _ = self.add_door(2, 1, 1, color='blue', locked=True)
        # Add a key to unlock the door
        self.add_object(2, 1, 'key', door.color)


        door, _ = self.add_door(2, 2, 2, color='purple', locked = False)
        # Add a key to unlock the door
        self.add_object(2, 2, 'key', door.color)


        door, _ = self.add_door(1, 2, 2,  color= 'green', locked = False)
        door, _ = self.add_door(0, 2, 3,  color= 'yellow', locked = False)
        door, _ = self.add_door(0, 1, 0,  color= 'red', locked = False)








        if self.distractors:
            self.add_distractors(num_distractors=4)

        self.place_agent(0, 0)

        self.instrs = PickupInstr(ObjDesc(obj.type, obj.color))





class Level_UnlockPickup9FixRoom(RoomGridLevel):
    """
    Unlock a door, then pick up a box in another room
    """

    def __init__(self, distractors=False, seed=None):
        self.distractors = distractors

        room_size = 7
        super().__init__(
            num_rows=3,
            num_cols=3,
            room_size=room_size,
            max_steps=8*room_size**2,
            seed=seed
        )

    def gen_mission(self):
        # Add a random object to the room on the right
        obj, _ = self.add_object(1, 1, color= 'yellow', kind="ball")
        # Make sure the two rooms are directly connected by a locked door


        door, _ = self.add_door(0, 0, 0, color= 'red', locked=True)
        # Add a key to unlock the door
        self.add_object(0, 0, 'key', door.color)


        door, _ = self.add_door(1, 0, 0, color='yellow', locked=True)
        # Add a key to unlock the door
        self.add_object(1, 0, 'key', door.color)


        door, _ = self.add_door(2, 0, 1, color='grey', locked=True)
        # Add a key to unlock the door
        self.add_object(2, 0, 'key', door.color)


        door, _ = self.add_door(2, 1, 1, color='blue', locked=True)
        # Add a key to unlock the door
        self.add_object(2, 1, 'key', door.color)


        door, _ = self.add_door(2, 2, 2, color='purple', locked = False)
        # Add a key to unlock the door
        self.add_object(2, 2, 'key', door.color)


        door, _ = self.add_door(1, 2, 2,  color= 'green', locked = True)
        self.add_object(1, 2, 'key', door.color)

        door, _ = self.add_door(0, 2, 3,  color= 'yellow', locked = False)
        door, _ = self.add_door(0, 1, 0,  color= 'red', locked = False)








        if self.distractors:
            self.add_distractors(num_distractors=4)

        self.place_agent(0, 0)

        self.instrs = PickupInstr(ObjDesc(obj.type, obj.color))





class Level_UnlockPickup9Room(RoomGridLevel):
    """
    Unlock a door, then pick up a box in another room
    """

    def __init__(self, distractors=False, seed=None):
        self.distractors = distractors

        room_size = 7
        super().__init__(
            num_rows=3,
            num_cols=3,
            room_size=room_size,
            max_steps=8*room_size**2,
            seed=seed
        )

    def gen_mission(self):
        # Add a random object to the room on the right
        obj, _ = self.add_object(1, 1, color= 'yellow', kind="ball")
        # Make sure the two rooms are directly connected by a locked door


        door, _ = self.add_door(0, 0, 0, locked=False)
        # Add a key to unlock the door
        # self.add_object(0, 0, 'key', door.color)


        door, _ = self.add_door(1, 0, 0,  locked=False)
        # Add a key to unlock the door
        # self.add_object(1, 0, 'key', door.color)


        door, _ = self.add_door(2, 0, 1,  locked=True)
        # Add a key to unlock the door
        self.add_object(2, 0, 'key', door.color)


        door, _ = self.add_door(2, 1, 1, locked=True)
        # Add a key to unlock the door
        self.add_object(2, 1, 'key', door.color)


        door, _ = self.add_door(2, 2, 2, locked = False)
        # Add a key to unlock the door
        self.add_object(2, 2, 'key', door.color)


        door, _ = self.add_door(1, 2, 2,  locked = False)
        # self.add_object(1, 2, 'key', door.color)

        door, _ = self.add_door(0, 2, 3, locked = False)
        door, _ = self.add_door(0, 1, 0, locked = False)








        if self.distractors:
            self.add_distractors(num_distractors=4)

        self.place_agent(0, 0)

        self.instrs = PickupInstr(ObjDesc(obj.type, obj.color))


class Level_UnlockPickup16Room(RoomGridLevel):
    """
    Unlock a door, then pick up a box in another room
    """

    def __init__(self, distractors=False, seed=None):
        self.distractors = distractors

        room_size = 7
        super().__init__(
            num_rows=4,
            num_cols=4,
            room_size=room_size,
            max_steps=8*room_size**2,
            seed=seed
        )

    def gen_mission(self):
        # Add a random object to the room on the right

          #######  3   #######
         #2 door  position 0#
        ########  1  #######


        obj, _ = self.add_object(2, 3, color= 'yellow', kind="ball")
        # Make sure the two rooms are directly connected by a locked door


        door, _ = self.add_door(0, 0, 0, color= 'red', locked=True)
        # Add a key to unlock the door
        self.add_object(0, 0, 'key', door.color)


        door, _ = self.add_door(1, 0, 1, color='yellow', locked=True)
        # Add a key to unlock the door
        self.add_object(1, 0, 'key', door.color)


        door, _ = self.add_door(2, 0, 1, color='grey', locked=True)
        # Add a key to unlock the door
        self.add_object(2, 1, 'key', door.color)


        door, _ = self.add_door(2, 1, 1, color='blue', locked=True)
        # Add a key to unlock the door
        self.add_object(2, 2, 'key', door.color)


        door, _ = self.add_door(2, 2, 2, color='purple', locked = False)
        # Add a key to unlock the door
        self.add_object(2, 2, 'key', door.color)


        door, _ = self.add_door(0, 2, 1,  color= 'green', locked = True)
        self.add_object(0, 2, 'key', door.color)

        door, _ = self.add_door(0, 2, 3,  color= 'yellow', locked = False)
        door, _ = self.add_door(0, 1, 0,  color= 'red', locked = False)

        door, _ = self.add_door(0, 3, 0,  color= 'red', locked = False)

        door, _ = self.add_door(1, 3, 3,  color= 'purple', locked = False)


        door, _ = self.add_door(2, 0, 0,  color= 'yellow', locked = False)

        door, _ = self.add_door(3, 0, 1,  color= 'blue', locked = False)

        door, _ = self.add_door(3, 1, 1,  color= 'red', locked = False)

        door, _ = self.add_door(3, 2, 1,  color= 'red', locked = False)

        door, _ = self.add_door(3, 3, 2,  color= 'grey', locked = False)


        if self.distractors:
            self.add_distractors(num_distractors=4)

        self.place_agent(0, 0)

        self.instrs = PickupInstr(ObjDesc(obj.type, obj.color))




class Level_ActionObjDoorOne(RoomGridLevel):
    """
    [pick up an object] or
    [go to an object or door] or
    [open a door]
    (in the current room)
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=7,
            seed=seed
        )

    def gen_mission(self):
        if random.random() < 0.5:
            objs = self.add_distractors(1, 1, num_distractors=1)
        else:
            objs = []
            door, _ = self.add_door(1, 1, locked=False)
            objs.append(door)

        self.place_agent(1, 1)

        obj = self._rand_elem(objs)
        desc = ObjDesc(obj.type, obj.color)

        if obj.type == 'door':
            self.instrs = OpenInstr(desc)
        else:
            self.instrs = PickupInstr(desc)



class Level_UnlockLocal(RoomGridLevel):
    """
    Fetch a key and unlock a door
    (in the current room)
    """

    def __init__(self, distractors=False, seed=None):
        self.distractors = distractors
        super().__init__(seed=seed)

    def gen_mission(self):
        door, _ = self.add_door(1, 1, locked=True)
        self.add_object(1, 1, 'key', door.color)
        if self.distractors:
            self.add_distractors(1, 1, num_distractors=3)
        self.place_agent(1, 1)

        self.instrs = OpenInstr(ObjDesc(door.type))


class Level_UnlockLocalDist(Level_UnlockLocal):
    """
    Fetch a key and unlock a door
    (in the current room, with distractors)
    """

    def __init__(self, seed=None):
        super().__init__(distractors=True, seed=seed)


class Level_KeyInBox(RoomGridLevel):
    """
    Unlock a door. Key is in a box (in the current room).
    """

    def __init__(self, seed=None):
        super().__init__(
            seed=seed
        )

    def gen_mission(self):
        door, _ = self.add_door(1, 1, locked=True)

        # Put the key in the box, then place the box in the room
        key = Key(door.color)
        box = Box(self._rand_color(), key)
        self.place_in_room(1, 1, box)

        self.place_agent(1, 1)

        self.instrs = OpenInstr(ObjDesc(door.type))


class Level_UnlockPickup(RoomGridLevel):
    """
    Unlock a door, then pick up a box in another room
    """

    def __init__(self, distractors=False, seed=None):
        self.distractors = distractors

        room_size = 7
        super().__init__(
            num_rows=1,
            num_cols=2,
            room_size=room_size,
            max_steps=8*room_size**2,
            seed=seed
        )

    def gen_mission(self):
        # Add a random object to the room on the right
        obj, _ = self.add_object(1, 0, kind="box")
        # Make sure the two rooms are directly connected by a locked door
        door, _ = self.add_door(0, 0, 0, locked=True)
        # Add a key to unlock the door
        self.add_object(0, 0, 'key', door.color)
        if self.distractors:
            self.add_distractors(num_distractors=4)

        self.place_agent(0, 0)

        self.instrs = PickupInstr(ObjDesc(obj.type, obj.color))


class Level_UnlockPickup2Room(RoomGridLevel):
    """
    Unlock a door, then pick up a box in another room
    """

    def __init__(self, distractors=False, seed=None):
        self.distractors = distractors

        room_size = 7
        super().__init__(
            num_rows=1,
            num_cols=3,
            room_size=room_size,
            max_steps=8*room_size**2,
            seed=seed
        )

    def gen_mission(self):
        # Add a random object to the room on the right
        obj, _ = self.add_object(1, 0, kind="box", color='purple')
        # Make sure the two rooms are directly connected by a locked door
        door, _ = self.add_door(0, 0, 0, locked=True)
        # Add a key to unlock the door
        self.add_object(0, 0, 'key', door.color)

        door, _ = self.add_door(1, 0, 0, locked=True)
        # Add a key to unlock the door
        self.add_object(1, 0, 'key', door.color)

        if self.distractors:
            self.add_distractors(num_distractors=4)

        self.place_agent(0, 0)

        self.instrs = PickupInstr(ObjDesc(obj.type, obj.color))


class Level_UnlockPickup3Room(RoomGridLevel):
    """
    Unlock a door, then pick up a box in another room
    """

    def __init__(self, distractors=False, seed=None):
        self.distractors = distractors

        room_size = 7
        super().__init__(
            num_rows=1,
            num_cols=3,
            room_size=room_size,
            max_steps=8*room_size**2,
            seed=seed
        )

    def gen_mission(self):
        # Add a random object to the room on the right
        obj, _ = self.add_object(2, 0, kind="box", color='purple')
        # Make sure the two rooms are directly connected by a locked door
        door, _ = self.add_door(0, 0, 0, locked=True)
        # Add a key to unlock the door
        self.add_object(0, 0, 'key', door.color)

        door, _ = self.add_door(1, 0, 0, locked=True)
        # Add a key to unlock the door
        self.add_object(1, 0, 'key', door.color)

        if self.distractors:
            self.add_distractors(num_distractors=4)

        self.place_agent(0, 0)

        self.instrs = PickupInstr(ObjDesc(obj.type, obj.color))


class Level_ShowRoom(RoomGridLevel):
    """
    Unlock a door, then pick up a box in another room
    """

    def __init__(self, distractors=False, seed=None):
        self.distractors = distractors

        room_size = 9
        super().__init__(
            num_rows=2,
            num_cols=3,
            room_size=room_size,
            max_steps=8*room_size**2,
            seed=seed
        )

    def gen_mission(self):
        # Add a random object to the room on the right
        obj, _ = self.add_object(2, 0, color='green', kind="box")
        # Make sure the two rooms are directly connected by a locked door
        door, _ = self.add_door(0, 0, 0, color= 'yellow', locked=True)
        # Add a key to unlock the door
        self.add_object(0, 0, 'key', door.color)

        door, _ = self.add_door(0, 0, 1, color='blue', locked=True)
        # Add a key to unlock the door
        # self.add_object(1, 0, 'key', door.color)

        if self.distractors:
            self.add_distractors(num_distractors=4)

        self.place_agent(0, 0)

        self.instrs = PickupInstr(ObjDesc(obj.type, obj.color))


class Level_UnlockPickupDist(Level_UnlockPickup):
    """
    Unlock a door, then pick up an object in another room
    (with distractors)
    """

    def __init__(self, seed=None):
        super().__init__(distractors=True, seed=seed)


class Level_BlockedUnlockPickup(RoomGridLevel):
    """
    Unlock a door blocked by a ball, then pick up a box
    in another room
    """

    def __init__(self, seed=None):
        room_size = 7
        super().__init__(
            num_rows=1,
            num_cols=2,
            room_size=room_size,
            max_steps=16*room_size**2,
            seed=seed
        )

    def gen_mission(self):
        # Add a box to the room on the right
        obj, _ = self.add_object(1, 0, kind="box")
        # Make sure the two rooms are directly connected by a locked door
        door, pos = self.add_door(0, 0, 0, locked=True)
        # Block the door with a ball
        color = self._rand_color()
        self.grid.set(pos[0]-1, pos[1], Ball(color))
        # Add a key to unlock the door
        self.add_object(0, 0, 'key', door.color)

        self.place_agent(0, 0)

        self.instrs = PickupInstr(ObjDesc(obj.type))


class Level_UnlockToUnlock(RoomGridLevel):
    """
    Unlock a door A that requires to unlock a door B before
    """

    def __init__(self, seed=None):
        room_size = 7
        super().__init__(
            num_rows=1,
            num_cols=3,
            room_size=room_size,
            max_steps=30*room_size**2,
            seed=seed
        )

    def gen_mission(self):
        colors = self._rand_subset(COLOR_NAMES, 2)

        # Add a door of color A connecting left and middle room
        self.add_door(0, 0, door_idx=0, color=colors[0], locked=True)

        # Add a key of color A in the room on the right
        self.add_object(2, 0, kind="key", color=colors[0])

        # Add a door of color B connecting middle and right room
        self.add_door(1, 0, door_idx=0, color=colors[1], locked=True)

        # Add a key of color B in the middle room
        self.add_object(1, 0, kind="key", color=colors[1])

        obj, _ = self.add_object(0, 0, kind="ball")

        self.place_agent(1, 0)

        self.instrs = PickupInstr(ObjDesc(obj.type))

###################################################################################################################
class Level_PickupDist(RoomGridLevel):
    """
    Pick up an object
    The object to pick up is given by its type only, or
    by its color, or by its type and color.
    (in the current room, with distractors)
    """

    def __init__(self, debug=False, seed=None):
        self.debug = debug
        super().__init__(
            num_rows = 1,
            num_cols = 1,
            room_size=7,
            seed=seed
        )

    def gen_mission(self):
        # Add 5 random objects in the room
        # obj = ('key','red')
        # room_i = self._rand_int(0, self.num_cols)
        #
        # room_j = self._rand_int(0, self.num_rows)
        #
        # dist, pos = self.add_object(room_i, room_j, *obj)
        #
        #
        #
        # objs = []
        # objs.append(dist)
        # # self.place_agent(0, 0)
        #
        self.place_agent()



        objs = self.add_distractors(num_distractors=5)

        obj = self._rand_elem(objs)


        type = obj.type
        color = obj.color

        # select_by = self._rand_elem(["type", "color", "both"])
        # if select_by == "color":
        #     type = None
        # elif select_by == "type":
        #     color = None

        self.instrs = PickupInstr(ObjDesc(type, color), strict=self.debug)

####################################################################################################################
class Level_PickupDistDebug(Level_PickupDist):
    """
    Same as PickupDist but the level stops when any object is picked
    """

    def __init__(self, seed=None):
        super().__init__(
            debug=True,
            seed=seed
        )


class Level_PickupAbove(RoomGridLevel):
    """
    Pick up an object (in the room above)
    This task requires to use the compass to be solved effectively.
    """

    def __init__(self, seed=None):
        room_size = 6
        super().__init__(
            room_size=room_size,
            max_steps=8*room_size**2,
            seed=seed
        )

    def gen_mission(self):
        # Add a random object to the top-middle room
        obj, pos = self.add_object(1, 0)
        # Make sure the two rooms are directly connected
        self.add_door(1, 1, 3, locked=False)
        self.place_agent(1, 1)
        self.connect_all()

        self.instrs = PickupInstr(ObjDesc(obj.type, obj.color))


class Level_OpenTwoDoors(RoomGridLevel):
    """
    Open door X, then open door Y
    The two doors are facing opposite directions, so that the agent
    Can't see whether the door behind him is open.
    This task requires memory (recurrent policy) to be solved effectively.
    """

    def __init__(self,
        first_color=None,
        second_color=None,
        strict=False,
        seed=None
    ):
        self.first_color = first_color
        self.second_color = second_color
        self.strict = strict

        room_size = 6
        super().__init__(
            room_size=room_size,
            max_steps=20*room_size**2,
            seed=seed
        )

    def gen_mission(self):
        colors = self._rand_subset(COLOR_NAMES, 2)

        first_color = self.first_color
        if first_color is None:
            first_color = colors[0]
        second_color = self.second_color
        if second_color is None:
            second_color = colors[1]

        door1, _ = self.add_door(1, 1, 2, color=first_color, locked=False)
        door2, _ = self.add_door(1, 1, 0, color=second_color, locked=False)

        self.place_agent(1, 1)

        self.instrs = BeforeInstr(
            OpenInstr(ObjDesc(door1.type, door1.color), strict=self.strict),
            OpenInstr(ObjDesc(door2.type, door2.color))
        )


class Level_OpenTwoDoorsDebug(Level_OpenTwoDoors):
    """
    Same as OpenTwoDoors but the level stops when the second door is opened
    """

    def __init__(self,
        first_color=None,
        second_color=None,
        seed=None
    ):
        super().__init__(
            first_color,
            second_color,
            strict=True,
            seed=seed
        )


class Level_OpenRedBlueDoors(Level_OpenTwoDoors):
    """
    Open red door, then open blue door
    The two doors are facing opposite directions, so that the agent
    Can't see whether the door behind him is open.
    This task requires memory (recurrent policy) to be solved effectively.
    """

    def __init__(self, seed=None):
        super().__init__(
            first_color="red",
            second_color="blue",
            seed=seed
        )


class Level_OpenRedBlueDoorsDebug(Level_OpenTwoDoorsDebug):
    """
    Same as OpenRedBlueDoors but the level stops when the blue door is opened
    """

    def __init__(self, seed=None):
        super().__init__(
            first_color="red",
            second_color="blue",
            seed=seed
        )


class Level_FindObjS5(RoomGridLevel):
    """
    Pick up an object (in a random room)
    Rooms have a size of 5
    This level requires potentially exhaustive exploration
    """

    def __init__(self, room_size=5, seed=None):
        super().__init__(
            room_size=room_size,
            max_steps=20*room_size**2,
            seed=seed
        )

    def gen_mission(self):
        # Add a random object to a random room
        i = self._rand_int(0, self.num_rows)
        j = self._rand_int(0, self.num_cols)
        obj, _ = self.add_object(i, j)
        self.place_agent(1, 1)
        self.connect_all()

        self.instrs = PickupInstr(ObjDesc(obj.type))


class Level_FindObjS6(Level_FindObjS5):
    """
    Same as the FindObjS5 level, but rooms have a size of 6
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=6,
            seed=seed
        )


class Level_FindObjS7(Level_FindObjS5):
    """
    Same as the FindObjS5 level, but rooms have a size of 7
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=7,
            seed=seed
        )


class KeyCorridor(RoomGridLevel):
    """
    A ball is behind a locked door, the key is placed in a
    random room.
    """

    def __init__(
        self,
        num_rows=3,
        obj_type="ball",
        room_size=6,
        seed=None
    ):
        self.obj_type = obj_type

        super().__init__(
            room_size=room_size,
            num_rows=num_rows,
            max_steps=30*room_size**2,
            seed=seed,
        )

    def gen_mission(self):
        # Connect the middle column rooms into a hallway
        for j in range(1, self.num_rows):
            self.remove_wall(1, j, 3)

        # Add a locked door on the bottom right
        # Add an object behind the locked door
        room_idx = self._rand_int(0, self.num_rows)
        door, _ = self.add_door(2, room_idx, 2, locked=True)
        obj, _ = self.add_object(2, room_idx, kind=self.obj_type)

        # Add a key in a random room on the left side
        self.add_object(0, self._rand_int(0, self.num_rows), 'key', door.color)

        # Place the agent in the middle
        self.place_agent(1, self.num_rows // 2)

        # Make sure all rooms are accessible
        self.connect_all()

        self.instrs = PickupInstr(ObjDesc(obj.type))


class Level_KeyCorridorS3R1(KeyCorridor):
    def __init__(self, seed=None):
        super().__init__(
            room_size=3,
            num_rows=1,
            seed=seed
        )

class Level_KeyCorridorS3R2(KeyCorridor):
    def __init__(self, seed=None):
        super().__init__(
            room_size=3,
            num_rows=2,
            seed=seed
        )

class Level_KeyCorridorS3R3(KeyCorridor):
    def __init__(self, seed=None):
        super().__init__(
            room_size=3,
            num_rows=3,
            seed=seed
        )

class Level_KeyCorridorS4R3(KeyCorridor):
    def __init__(self, seed=None):
        super().__init__(
            room_size=4,
            num_rows=3,
            seed=seed
        )

class Level_KeyCorridorS5R3(KeyCorridor):
    def __init__(self, seed=None):
        super().__init__(
            room_size=5,
            num_rows=3,
            seed=seed
        )

class Level_KeyCorridorS6R3(KeyCorridor):
    def __init__(self, seed=None):
        super().__init__(
            room_size=6,
            num_rows=3,
            seed=seed
        )

class Level_1RoomS8(RoomGridLevel):
    """
    Pick up the ball
    Rooms have a size of 8
    """

    def __init__(self, room_size=8, seed=None):
        super().__init__(
            room_size=room_size,
            num_rows=1,
            num_cols=1,
            seed=seed
        )

    def gen_mission(self):
        obj, _ = self.add_object(0, 0, kind="ball")
        self.place_agent()
        self.instrs = PickupInstr(ObjDesc(obj.type))


class Level_1RoomS12(Level_1RoomS8):
    """
    Pick up the ball
    Rooms have a size of 12
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=12,
            seed=seed
        )


class Level_1RoomS16(Level_1RoomS8):
    """
    Pick up the ball
    Rooms have a size of 16
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=16,
            seed=seed
        )


class Level_1RoomS20(Level_1RoomS8):
    """
    Pick up the ball
    Rooms have a size of 20
    """

    def __init__(self, seed=None):
        super().__init__(
            room_size=20,
            seed=seed
        )


class PutNext(RoomGridLevel):
    """
    Task of the form: move the A next to the B and the C next to the D.
    This task is structured to have a very large number of possible
    instructions.
    """

    def __init__(
        self,
        room_size,
        objs_per_room,
        start_carrying=False,
        seed=None
    ):
        assert room_size >= 4
        assert objs_per_room <= 9
        self.objs_per_room = objs_per_room
        self.start_carrying = start_carrying

        super().__init__(
            num_rows=1,
            num_cols=2,
            room_size=room_size,
            max_steps=8*room_size**2,
            seed=seed
        )

    def gen_mission(self):
        self.place_agent(0, 0)

        # Add objects to both the left and right rooms
        # so that we know that we have two non-adjacent set of objects
        objs_l = self.add_distractors(0, 0, self.objs_per_room)
        objs_r = self.add_distractors(1, 0, self.objs_per_room)

        # Remove the wall between the two rooms
        self.remove_wall(0, 0, 0)

        # Select objects from both subsets
        a = self._rand_elem(objs_l)
        b = self._rand_elem(objs_r)

        # Randomly flip the object to be moved
        if self._rand_bool():
            t = a
            a = b
            b = t

        self.obj_a = a

        self.instrs = PutNextInstr(
            ObjDesc(a.type, a.color),
            ObjDesc(b.type, b.color)
        )

    def reset(self, **kwargs):
        obs = super().reset(**kwargs)

        # If the agent starts off carrying the object
        if self.start_carrying:
            self.grid.set(*self.obj_a.init_pos, None)
            self.carrying = self.obj_a

        return obs


class Level_PutNextS4N1(PutNext):
    def __init__(self, seed=None):
        super().__init__(
            room_size=4,
            objs_per_room=1,
            seed=seed
        )


class Level_PutNextS5N1(PutNext):
    def __init__(self, seed=None):
        super().__init__(
            room_size=5,
            objs_per_room=1,
            seed=seed
        )


class Level_PutNextS5N2(PutNext):
    def __init__(self, seed=None):
        super().__init__(
            room_size=5,
            objs_per_room=2,
            seed=seed
        )


class Level_PutNextS6N3(PutNext):
    def __init__(self, seed=None):
        super().__init__(
            room_size=6,
            objs_per_room=3,
            seed=seed
        )


class Level_PutNextS7N4(PutNext):
    def __init__(self, seed=None):
        super().__init__(
            room_size=7,
            objs_per_room=4,
            seed=seed
        )


class Level_PutNextS5N2Carrying(PutNext):
    def __init__(self, seed=None):
        super().__init__(
            room_size=5,
            objs_per_room=2,
            start_carrying=True,
            seed=seed
        )


class Level_PutNextS6N3Carrying(PutNext):
    def __init__(self, seed=None):
        super().__init__(
            room_size=6,
            objs_per_room=3,
            start_carrying=True,
            seed=seed
        )


class Level_PutNextS7N4Carrying(PutNext):
    def __init__(self, seed=None):
        super().__init__(
            room_size=7,
            objs_per_room=4,
            start_carrying=True,
            seed=seed
        )


class MoveTwoAcross(RoomGridLevel):
    """
    Task of the form: move the A next to the B and the C next to the D.
    This task is structured to have a very large number of possible
    instructions.
    """

    def __init__(
        self,
        room_size,
        objs_per_room,
        seed=None
    ):
        assert objs_per_room <= 9
        self.objs_per_room = objs_per_room

        super().__init__(
            num_rows=1,
            num_cols=2,
            room_size=room_size,
            max_steps=16*room_size**2,
            seed=seed
        )

    def gen_mission(self):
        self.place_agent(0, 0)

        # Add objects to both the left and right rooms
        # so that we know that we have two non-adjacent set of objects
        objs_l = self.add_distractors(0, 0, self.objs_per_room)
        objs_r = self.add_distractors(1, 0, self.objs_per_room)

        # Remove the wall between the two rooms
        self.remove_wall(0, 0, 0)

        # Select objects from both subsets
        objs_l = self._rand_subset(objs_l, 2)
        objs_r = self._rand_subset(objs_r, 2)
        a = objs_l[0]
        b = objs_r[0]
        c = objs_r[1]
        d = objs_l[1]

        self.instrs = BeforeInstr(
            PutNextInstr(ObjDesc(a.type, a.color), ObjDesc(b.type, b.color)),
            PutNextInstr(ObjDesc(c.type, c.color), ObjDesc(d.type, d.color))
        )


class Level_MoveTwoAcrossS5N2(MoveTwoAcross):
    def __init__(self, seed=None):
        super().__init__(
            room_size=5,
            objs_per_room=2,
            seed=seed
        )


class Level_MoveTwoAcrossS8N9(MoveTwoAcross):
    def __init__(self, seed=None):
        super().__init__(
            room_size=8,
            objs_per_room=9,
            seed=seed
        )


class OpenDoorsOrder(RoomGridLevel):
    """
    Open one or two doors in the order specified.
    """

    def __init__(
        self,
        num_doors,
        debug=False,
        seed=None
    ):
        assert num_doors >= 2
        self.num_doors = num_doors
        self.debug = debug

        room_size = 6
        super().__init__(
            room_size=room_size,
            max_steps=20*room_size**2,
            seed=seed
        )

    def gen_mission(self):
        colors = self._rand_subset(COLOR_NAMES, self.num_doors)
        doors = []
        for i in range(self.num_doors):
            door, _ = self.add_door(1, 1, color=colors[i], locked=False)
            doors.append(door)
        self.place_agent(1, 1)

        door1, door2 = self._rand_subset(doors, 2)
        desc1 = ObjDesc(door1.type, door1.color)
        desc2 = ObjDesc(door2.type, door2.color)

        mode = self._rand_int(0, 3)
        if mode == 0:
            self.instrs = OpenInstr(desc1, strict=self.debug)
        elif mode == 1:
            self.instrs = BeforeInstr(OpenInstr(desc1, strict=self.debug), OpenInstr(desc2, strict=self.debug))
        elif mode == 2:
            self.instrs = AfterInstr(OpenInstr(desc1, strict=self.debug), OpenInstr(desc2, strict=self.debug))
        else:
            assert False

class Level_OpenDoorsOrderN2(OpenDoorsOrder):
    def __init__(self, seed=None):
        super().__init__(
            num_doors=2,
            seed=seed
        )


class Level_OpenDoorsOrderN4(OpenDoorsOrder):
    def __init__(self, seed=None):
        super().__init__(
            num_doors=4,
            seed=seed
        )


class Level_OpenDoorsOrderN2Debug(OpenDoorsOrder):
    def __init__(self, seed=None):
        super().__init__(
            num_doors=2,
            debug=True,
            seed=seed
        )


class Level_OpenDoorsOrderN4Debug(OpenDoorsOrder):
    def __init__(self, seed=None):
        super().__init__(
            num_doors=4,
            debug=True,
            seed=seed
        )

for name, level in list(globals().items()):
    if name.startswith('Level_'):
        level.is_bonus = True

# Register the levels in this file
register_levels(__name__, globals())
