"""
Experiment demonstrating task composition
"""
import torch
import numpy as np
import itertools

from dqn.libraries import *
from envs.env import *

    
if __name__ == '__main__':
    obj_color = 'blue'
    obj_type = 'box'
    path = './dqn/models'
    
    print('Loading ...')
    env = make_env()
    if obj_type:
        goals = list(itertools.product(OBJ_COLORS,[obj_type],env.rooms))
    if obj_color:
        goals = list(itertools.product([obj_color],OBJ_TYPES,env.rooms))
    env = make_env(goals=None)
    # color = load('{0}/{1}.dqn'.format(path,obj_color), env)
    obj = load('{0}/{1}.dqn'.format(path,'object'), env)

    print('Composing ...')
    # model = AND(color,obj)
    model = obj
    
    print('Visualizing ...')        
    max_episodes = 50000
    max_trajectory = 20
    for episode in range(max_episodes):
        obs = env.reset()
        goal = select_goal(model,obs)
        
        for _ in range(max_trajectory):
            env.render()
            action = select_action(model,obs,goal)
            obs, reward, done, _ = env.step(action)
            
            if done or env.window.closed:
                print(reward)
                break

        if env.window.closed:
            break
