import numpy as np
import torch
import os
import itertools
import sys

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

path = './dqn/models'

if __name__ == '__main__':    
    task = {'obj_color': None, 'obj_type': None, 'room': None}
    name = ""
    name = name + task['obj_color'] if task['obj_color'] else name
    name = name + task['obj_type'] if task['obj_type'] else name
    name = name + task['room'] if task['room'] else name
    name = name if name!="" else "object"
    model_path = "{0}/{1}.dqn".format(path,name)  

    # Initialising
    env = make_env()
    if task['obj_type']:
        goals = list(itertools.product(OBJ_COLORS,[task['obj_type']],env.rooms))
    elif task['obj_color']:
        goals = list(itertools.product([task['obj_color']],OBJ_TYPES,env.rooms))
    elif task['room']:
        goals = list(itertools.product(OBJ_COLORS,OBJ_TYPES,[task['room']]))
    else:
        goals = env.all_goals
    env = make_env(goals=goals)

    # Training
    max_timesteps_rl = int(1e7)
    model, steps = train(env, type_="vf", max_timesteps=max_timesteps_rl, path=model_path, save_model=True)
    save(model_path, model)
    

