"""Commands A1 robot to raise and lower its legs so it crouches and stands up.

Can be run in sim by setting --real_robot=False.
"""

import inspect
import os

from absl import app
from absl import flags
from absl import logging
import numpy as np
import time
from tqdm import tqdm

current_dir = os.path.dirname(os.path.abspath(__file__))
dependencies_dir = os.path.join(os.path.join(current_dir, "../envs"), "locomotion_simulation")
os.sys.path.insert(0, dependencies_dir)

import pybullet  # pytype:disable=import-error
import pybullet_data
from pybullet_utils import bullet_client

from locomotion.robots import a1_robot
from locomotion.robots import robot_config

FREQ = 0.5


def main(_):

  logging.info(
      "WARNING: this code executes low-level controller on the robot.")
  logging.info("Make sure the robot is hang on rack before proceeding.")
  input("Press enter to continue...")

  # Construct sim env and real robot
  action_repeat = 10
  time_step = 0.001
  p = bullet_client.BulletClient(connection_mode=pybullet.DIRECT)
  p.setAdditionalSearchPath(pybullet_data.getDataPath())
  robot = a1_robot.A1Robot(pybullet_client=p, 
                           action_repeat=action_repeat,
                           time_step=time_step,
                           reset_time=-1)
  print(f'time_step: {time_step}, action_repeat: {action_repeat}, Step costs: {action_repeat*time_step}')

  # Move the motors slowly to initial position
  robot.ReceiveObservation()
  current_motor_angle = np.array(robot.GetMotorAngles())
  desired_motor_angle = np.array([0., 0.9, -1.8] * 4)
  print(f'\nCurrent Angle: {current_motor_angle}\nTarget Angle: {desired_motor_angle}')
  input("Press enter to stand...")
  for t in tqdm(range(300)):
    t0 = time.time()
    blend_ratio = np.minimum(t / 200., 1)
    action = (1 - blend_ratio
              ) * current_motor_angle + blend_ratio * desired_motor_angle
    robot.Step(action, robot_config.MotorControlMode.POSITION)
  
  # Move the legs in a sinusoidal curve
  input("Press enter to move...")
  for t in tqdm(range(1000)):
    angle_hip = 0.9 + 0.2 * np.sin(2 * np.pi * FREQ * 0.01 * t)
    angle_calf = -2 * angle_hip
    action = np.array([0., angle_hip, angle_calf] * 4)
    robot.Step(action, robot_config.MotorControlMode.POSITION)

  robot.Terminate()


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