from domain.WorldGenerator import generate_world
from domain.levels.Task import make_object, _make_block, _make_item, make_door
from domain.levels.layout import MAPS
from domain.object_wrapper import ObjectWrapper
from domain.objects import Rect, Door, Block, Item
from domain.options import WalkDoor, ToggleDoor, WalkThroughDoor, WalkToItem, PickupItem, AttackItem, CraftItem, \
    OpenChest, WalkNorthDoor, WalkSouthDoor
import numpy as np
import random


def make_fixed_object(id, object_creator, pos, room):
    if id == 1:
        x = 1 if random.random() else 8
        pos[0] = x
        # pos[0] = 8

    object = object_creator(pos, room)
    if 'pickaxe' in object.type:
        object.set_look_at(pos[0], pos[1], pos[2] - 1, 45, 0)
    else:
        object.set_look_at(pos[0], pos[1], pos[2] - 3, 30, 0)
    object.id = id
    return object


def generate(seed, object_view=True, record=False, **kwargs):
    fixed_seed = 31
    layout = MAPS[fixed_seed]

    random.seed(seed)

    room1 = Rect(0, 0, 9, 6)
    room2 = Rect(0, 8, 9, 14)
    room3 = Rect(0, 16, 9, 22)
    room4 = Rect(0, 24, 9, 30)
    room5 = Rect(0, 32, 9, 39)

    door1, door2, door3, door4 = layout['door1'], layout['door2'], layout['door3'], layout['door4']
    gold_pos = layout['gold']

    rooms = [room1, room2, room3, room4]

    if random.random():
        p_room = room1
        rooms.pop(0)
    else:
        p_room = room2
        rooms.pop(1)

    random.shuffle(rooms)
    pickaxe = make_object(1, _make_item('diamond_pickaxe'), p_room)
    table = make_object(0, _make_block('crafting_table', False), rooms[0])
    chest = make_object(2, _make_block('chest', False), rooms[1])
    ore = make_object(3, _make_block('redstone_ore'), rooms[2])

    gold = Block(gold_pos[0], gold_pos[1], gold_pos[2], "gold_block", room5)
    gold.set_look_at(gold_pos[0], gold_pos[1], gold_pos[2] - 3, 30, 0)
    gold.id = 4

    objects = [table, pickaxe, chest, ore]  # leave gold out since env creates it!

    env = generate_world(seed=fixed_seed, palette='dungeon', objects=objects, wrap=object_view, record=record, **kwargs)
    objects.append(gold)

    door_1 = make_door(5, door1, room2, room1, layout['puzzles'][0])
    door_2 = make_door(6, door2, room3, room2, layout['puzzles'][1])
    door_3 = make_door(7, door3, room4, room3, layout['puzzles'][2])
    door_4 = make_door(8, door4, room5, room4, layout['puzzles'][3])

    doors = [door_1, door_2, door_3, door_4]

    if not object_view:
        return env

    env = ObjectWrapper(env, doors, objects)

    return env, doors, objects

if __name__ == '__main__':

    TASK_ID = 2
    seeds = [31, 33, 76, 82, 92]
    env = generate(seeds[TASK_ID], object_view=False, record=True)
    env.reset()