"""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__))
dependencies_dir = os.path.join(os.path.join(current_dir, "../envs"), "fine-tuning-locomotion")
os.sys.path.insert(0, dependencies_dir)
from motion_imitation.envs import env_builder

os.sys.path.insert(1, os.path.join(current_dir, "../"))
import utils 

# CONSTANTS
total_timesteps = 100
command_time = 0.01 # num_action_repeat * sim_time_step_s = 10 * 0.001
task_name = 'reset'


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_ftl_{task_name}-default/12345')
  results_path = agent_dir 
  cfg = utils.load_hydra_cfg(results_path)
  cfg.agent.params.obs_dim = 84
  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 = env_builder.build_env(
                        task_name,
                        mode="test",
                        enable_randomizer=False,
                        enable_rendering=False,
                        reset_at_current_position=False,
                        use_real_robot=FLAGS.real_robot,
                        realistic_sim=False)
  print(f"Config: {env._gym_config}")

  # Reset real robot
  print('\nResetting env...')
  obs = env.reset(
          initial_motor_angles=None,
          reset_duration=2.0)

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

  last_frame_time = time.time()
  for step in range(total_timesteps):
      # control fixed control time step (normal time step is 0.01s)
      time_spent = time.time() - last_frame_time
      last_frame_time = time.time()
      time_to_sleep = command_time - time_spent
      if time_to_sleep > 0:
        time.sleep(time_to_sleep)

      action = agent.act(obs, sample=False)
      next_obs, reward, done, info = env.step(action)

      # print(step, obs, action, reward, next_obs)
      obs = next_obs
  
  env._robot.Terminate()

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



