"""Apply the same action to the simulated and real A1 robot.


As a basic debug tool, this script allows you to execute the same action
(which you choose from the pybullet GUI) on the simulation and real robot
simultaneouly. Make sure to put the real robot on rack before testing.
"""

import os, sys
import time

import hydra

import numpy as np
import pybullet as p  # pytype: disable=import-error
from absl import app
from absl import flags
from absl import logging

current_dir = os.path.dirname(os.path.abspath(__file__))
os.sys.path.insert(1, os.path.join(current_dir, "../"))
import utils 
import envs
from envs.a1_env_ls import A1EnvLS

# Constants
total_timesteps = 10
command_time = 0.01 # num_action_repeat * sim_time_step_s = 10 * 0.001
task_name = 'sit'

# Parameters
flags.DEFINE_bool(
    "real_robot", True, "Whether to control a real robot or simulated.")
FLAGS = flags.FLAGS


def main(_):

  # Load agents 
  print("\nLoading agent...")
  agent_dir = os.path.join(current_dir, f'../outputs/a1-a1_ls_{task_name}-adv/12345')
  results_path = agent_dir 
  cfg = utils.load_hydra_cfg(results_path)
  cfg.agent.params.obs_dim = 19
  cfg.agent.params.action_dim = 12
  cfg.agent.params.action_range = [-1, 1]
  cfg.agent.params.device = "cpu"
  agent = hydra.utils.instantiate(cfg.agent)
  agent.load(agent_dir)

  # Construct real robot
  print("\nBuilding env...")
  env = A1EnvLS(task_name, real_robot=FLAGS.real_robot)
  print(f'\nGym Config: {env._env._gym_config}\n')
  # TODO: enable_action_filter

  # Reset real robot
  input('\nReset env?')
  obs = env.reset(reset_duration=3.0)
  print(f'\nInitial Obs: {obs}')

  # Run
  input(f"\nAre you ready for {total_timesteps} steps? ")
  print("\nGo!")

  last_frame_time = time.time()
  rewards, times = [], []
  for step in range(total_timesteps):

      action = agent.act(obs, sample=False)

      t0 = time.time() # real robot execution time
      next_obs, reward, done, info = env.step(action)
      t1 = time.time()

      times.append(t1-t0)
      rewards.append(reward)

      print(f'\nStep: {step} Reward:{rewards[-1]} Time:{times[-1]:.4f} \nObs: {obs[-12:]}\nAct: {action}\n')
      obs = next_obs
  
  print(f'\nLength: {step} Return: {np.sum(rewards)} Time: {np.sum(times)}')
  env._robot.Terminate()

if __name__ == "__main__":
  app.run(main)



