# AntCustom transferred to other angles
import subprocess
import os
import numpy as np
import h5py

get_env = {
    'HalfCheetah-v2': 'get_cheetah_env',
    'Walker2d-v2': 'get_walker_env',
    'Hopper-v2': 'get_hopper_env',
    'AntCustom-v2': 'get_ant_custom_env'
}
environ = os.environ.copy()
environ["CUDA_VISIBLE_DEVICES"] = "0"
path_prefix = os.environ['UDG_DATA_PATH']

env = 'AntCustom-v2'
l_env, _ = env.split('-')
l_env = l_env.lower()
u_env = l_env.upper()

model_code = 'c15'
model_path = '2023-04-29_08-29-20'
model_type = 'wasserstein'
model_step = 'final'
model_file = 'experience.h5'
pool_size = 1000000
rollout_length = 2
penalty_coeff = 7

# find the model with maximum reward
datasets_path = path_prefix+"url-data/data/"+model_type+"/"+env+"/"+model_path+"/"
reward_matrix = np.load(datasets_path+"reward_matrix.np", allow_pickle=True)
num_modes, num_angles = reward_matrix.shape
min_interval = 360 // num_angles
requested_angle = 0
requested_index = ((requested_angle + min_interval//2)//min_interval) % num_angles
mode = int(np.argmax(reward_matrix[:,requested_index]))
model_step = model_step + "_m{}".format(mode)
model_code = model_code + "_m{}_{}_{}".format(mode, rollout_length, penalty_coeff)

# relabel data
dataset_path = datasets_path+model_step+"/"+model_file
new_dataset_path = datasets_path+model_step+"/experience_angle_{}.h5".format(requested_angle)

# register env and write config file
dataset_id = l_env+"-"+model_code+"-"+str(requested_angle)

subprocess.run(["python", "examples/development/test_render.py", 
    "--config=examples.config.url."+dataset_id.replace('-','_'), "--gpus=1", "--trial-gpus=1", "--checkpoint-frequency=1"], env=environ)