import torch


def init_control(ctrl_obj, env, live_plot_obj, rec, params_general, random_actions_init,
					costs_tests, idx_test, num_repeat_actions=1):
	"""
	Init the gym env and memory. Define the lists containing the points for visualisations.
	Control the env with random actions for a number of steps.
	Args:
		ctrl_obj (control_objects.gp_mpc_controller.GpMpcController):
																Object containing the control functions.
																Used here to compute cost and add_memory.
		env (gym env): environment used to get the next observation given current observation and action
		live_plot_obj (object): object used for visualisation in real time of the control in a 2d graph.
								Must contain a function called update, that adds new points on the graph
		rec (object): object used to visualise the gym env in real-time.
						Must heve a function called capture_frame that updates the visualisation
		params_general (dict): see parameters.md for more information
		random_actions_init (int): number of random actions to apply for the initialization.
		costs_tests (numpy.array): contains the costs of all the tests realized, which is used to compute
									the mean cost over a large number of control. Parameter not relevant if num_tests=1.
									Dim=(num_runs, num_timesteps)
		idx_test (int): index of the test realized
		num_repeat_actions (int): number of consecutive actions that are constant.
									The GPs only store points at which the control change.
	"""
	# Lists used for plotting values
	obs_lst = []
	actions_lst = []
	rewards_lst = []
	env.reset()
	obs, reward, done, info = env.step(env.action_space.sample())
	obs_prev_ctrl = None
	action = None
	cost = None
	for idx_action in range(random_actions_init):
		if idx_action % num_repeat_actions == 0:
			if obs_prev_ctrl is not None and action is not None and cost is not None:
				ctrl_obj.add_memory(obs=obs_prev_ctrl, action=action, obs_new=obs,
									reward=-cost, check_storage=False)
			obs_prev_ctrl = obs
			action = env.action_space.sample()
		obs_new, reward, done, info = env.step(action)
		if params_general['render_env']:
			try: env.render()
			except: pass
		if params_general['save_render_env']:
			try: rec.capture_frame()
			except: pass

		cost, cost_var = ctrl_obj.compute_cost_unnormalized(obs_new, action)
		costs_tests[idx_test, idx_action] = cost
		obs_lst.append(obs)
		actions_lst.append(action)
		rewards_lst.append(-cost)

		if params_general['render_live_plots_2d']:
			live_plot_obj.update(obs=obs, action=action, cost=cost, info_dict=None)

		obs = obs_new
		# Necessary to store the last action in case that the parameter limit_action_change is set to True
		ctrl_obj.action_previous_iter = torch.Tensor(action)

	return ctrl_obj, env, live_plot_obj, rec, obs, action, cost, obs_prev_ctrl, \
			costs_tests, obs_lst, actions_lst, rewards_lst