import random
import time
import gym
import gym_minigrid.minigrid as minigrid
import networkx as nx
from networkx import grid_graph
import numpy as np
import torch as th
import gym_minigrid.minigrid as minigrid
import matplotlib.pyplot as plt
import argparse

from envs.multigrid import *
from envs.multigrid.adversarial import *
from envs.box2d import *
from envs.bipedalwalker import *

from envs.registration import make as gym_make
from envs.multigrid.adversarial import AdversarialEnv, GoalLastVariableBlocksAdversarialEnv
from add.diffusion_human_feedback.predictor import Tutor
from add.diffusion_human_feedback.generator import EnvGenerator
from util import create_parallel_env, DotDict, str2bool, seed

from arguments import parser
import time
import pyvirtualdisplay

if __name__ == '__main__':
	parser.add_argument(
		'--tutor_dir',
		type=str,
		default="../racing/minigrid_60/add/seed_1_cvar_03_guide_5/tutors/",
		help="path to the directory where tutor networks are saved"
	)
	parser.add_argument(
		'--tutor_model',
		type=str,
		default="model_030000.pt",
		help="name of model file"
	)
	args = parser.parse_args()
	args.env_name="BipedalWalker-Adversarial-v0"
	args.num_processes=1
	# if args.env_name.startswith("CarRacing"):
	# 	display = pyvirtualdisplay.Display(visible=0, size=(1400, 900), color_depth=24)
	# 	display.start()
	# env = gym.make("BipedalWalker-v3")
	env = BipedalWalkerAdversarialEnv()
	# env.reset_random()
	env.reset()
 
	env_params = np.load("../logs/bipedal/add/seed_2/env_params/61030.npy")
	env_params = env_params.astype(np.float32)
	for k in range(env_params.shape[0]):
		print(env_params[k])
		env_param = np.concatenate([env_params[k], [np.random.randint(1000)]])
		print(env_param.shape)
		env.reset_to_level(env_param)
		for i in range(5):
			state = env.render(mode='rgb_array')
			plt.imshow(env.viewer.get_array())
			plt.savefig(f"tt/bipedal_last_16/{k}.png", dpi=300)
	env.close()
