#The example to run the raibert controller in a Minitaur gym env.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import inspect
currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
parentdir = os.path.dirname(os.path.dirname(os.path.dirname(currentdir)))
os.sys.path.insert(0, parentdir)

import tf.compat.v1 as tf
from pybullet_envs.minitaur.envs import minitaur_raibert_controller
from pybullet_envs.minitaur.envs import minitaur_gym_env

flags = tf.app.flags
FLAGS = tf.app.flags.FLAGS

flags.DEFINE_float("motor_kp", 1.0, "The position gain of the motor.")
flags.DEFINE_float("motor_kd", 0.015, "The speed gain of the motor.")
flags.DEFINE_float("control_latency", 0.02, "The latency between sensor measurement and action"
                   " execution the robot.")
flags.DEFINE_string("log_path", None, "The directory to write the log file.")


def speed(t):
  max_speed = 0.35
  t1 = 3
  if t < t1:
    return t / t1 * max_speed
  else:
    return -max_speed


def main(argv):
  del argv
  try:
    env = minitaur_gym_env.MinitaurGymEnv(
        urdf_version=minitaur_gym_env.RAINBOW_DASH_V0_URDF_VERSION,
        control_time_step=0.006,
        action_repeat=6,
        pd_latency=0.003,
        control_latency=FLAGS.control_latency,
        motor_kp=FLAGS.motor_kp,
        motor_kd=FLAGS.motor_kd,
        remove_default_joint_damping=True,
        leg_model_enabled=False,
        render=True,
        on_rack=False,
        accurate_motor_model_enabled=True,
        log_path=FLAGS.log_path)
    env.reset()

    controller = minitaur_raibert_controller.MinitaurRaibertTrottingController(env.minitaur)

    tstart = env.minitaur.GetTimeSinceReset()
    for _ in range(1000):
      t = env.minitaur.GetTimeSinceReset() - tstart
      controller.behavior_parameters = (minitaur_raibert_controller.BehaviorParameters(
          desired_forward_speed=speed(t)))
      controller.update(t)
      env.step(controller.get_action())
  finally:
    env.close()


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