import pybullet_envs.minitaur.envs_v2.locomotion_gym_config
import pybullet_envs.minitaur.envs_v2.locomotion_gym_env
import pybullet_envs.minitaur.robots.mini_cheetah
import pybullet_data as pd

URDF_ROOT = pd.getDataPath()+"/urdf/"

ABDUCTION_P_GAIN = 100.0
ABDUCTION_D_GAIN = 1.0
HIP_P_GAIN = 30
HIP_D_GAIN = 2.0
KNEE_P_GAIN = 50
KNEE_D_GAIN = 2.0


UPPER_BOUND = 6.28318548203
LOWER_BOUND = -6.28318548203
SIM_TIME_STEP = 0.001
NUM_ACTION_REPEAT = 2
NUM_MOTORS = 12
NOISY_READING = True

locomotion_gym_config.SimulationParameters.sim_time_step_s = %SIM_TIME_STEP
locomotion_gym_config.SimulationParameters.num_action_repeat = %NUM_ACTION_REPEAT
locomotion_gym_config.SimulationParameters.enable_rendering = False
locomotion_gym_config.LocomotionGymConfig.simulation_parameters = @locomotion_gym_config.SimulationParameters()

Act0/locomotion_gym_config.ScalarField.name = "motor_angle_0"
Act0/locomotion_gym_config.ScalarField.upper_bound = %UPPER_BOUND
Act0/locomotion_gym_config.ScalarField.lower_bound = %LOWER_BOUND
Act1/locomotion_gym_config.ScalarField.name = "motor_angle_1"
Act1/locomotion_gym_config.ScalarField.upper_bound = %UPPER_BOUND
Act1/locomotion_gym_config.ScalarField.lower_bound = %LOWER_BOUND
Act2/locomotion_gym_config.ScalarField.name = "motor_angle_2"
Act2/locomotion_gym_config.ScalarField.upper_bound = %UPPER_BOUND
Act2/locomotion_gym_config.ScalarField.lower_bound = %LOWER_BOUND
Act3/locomotion_gym_config.ScalarField.name = "motor_angle_3"
Act3/locomotion_gym_config.ScalarField.upper_bound = %UPPER_BOUND
Act3/locomotion_gym_config.ScalarField.lower_bound = %LOWER_BOUND
Act4/locomotion_gym_config.ScalarField.name = "motor_angle_4"
Act4/locomotion_gym_config.ScalarField.upper_bound = %UPPER_BOUND
Act4/locomotion_gym_config.ScalarField.lower_bound = %LOWER_BOUND
Act5/locomotion_gym_config.ScalarField.name = "motor_angle_5"
Act5/locomotion_gym_config.ScalarField.upper_bound = %UPPER_BOUND
Act5/locomotion_gym_config.ScalarField.lower_bound = %LOWER_BOUND
Act6/locomotion_gym_config.ScalarField.name = "motor_angle_6"
Act6/locomotion_gym_config.ScalarField.upper_bound = %UPPER_BOUND
Act6/locomotion_gym_config.ScalarField.lower_bound = %LOWER_BOUND
Act7/locomotion_gym_config.ScalarField.name = "motor_angle_7"
Act7/locomotion_gym_config.ScalarField.upper_bound = %UPPER_BOUND
Act7/locomotion_gym_config.ScalarField.lower_bound = %LOWER_BOUND
Act8/locomotion_gym_config.ScalarField.name = "motor_angle_8"
Act8/locomotion_gym_config.ScalarField.upper_bound = %UPPER_BOUND
Act8/locomotion_gym_config.ScalarField.lower_bound = %LOWER_BOUND
Act9/locomotion_gym_config.ScalarField.name = "motor_angle_9"
Act9/locomotion_gym_config.ScalarField.upper_bound = %UPPER_BOUND
Act9/locomotion_gym_config.ScalarField.lower_bound = %LOWER_BOUND
Act10/locomotion_gym_config.ScalarField.name = "motor_angle_10"
Act10/locomotion_gym_config.ScalarField.upper_bound = %UPPER_BOUND
Act10/locomotion_gym_config.ScalarField.lower_bound = %LOWER_BOUND
Act11/locomotion_gym_config.ScalarField.name = "motor_angle_11"
Act11/locomotion_gym_config.ScalarField.upper_bound = %UPPER_BOUND
Act11/locomotion_gym_config.ScalarField.lower_bound = %LOWER_BOUND


locomotion_gym_config.LocomotionGymConfig.actions = [
    @Act0/locomotion_gym_config.ScalarField(),
    @Act1/locomotion_gym_config.ScalarField(),
    @Act2/locomotion_gym_config.ScalarField(),
    @Act3/locomotion_gym_config.ScalarField(),
    @Act4/locomotion_gym_config.ScalarField(),
    @Act5/locomotion_gym_config.ScalarField(),
    @Act6/locomotion_gym_config.ScalarField(),
    @Act7/locomotion_gym_config.ScalarField(),
    @Act8/locomotion_gym_config.ScalarField(),
    @Act9/locomotion_gym_config.ScalarField(),
    @Act10/locomotion_gym_config.ScalarField(),
    @Act11/locomotion_gym_config.ScalarField()]

locomotion_gym_env.LocomotionGymEnv.gym_config = @locomotion_gym_config.LocomotionGymConfig()

mini_cheetah.MiniCheetah.urdf_root = %URDF_ROOT
mini_cheetah.MiniCheetah.time_step = %SIM_TIME_STEP
mini_cheetah.MiniCheetah.action_repeat = %NUM_ACTION_REPEAT
mini_cheetah.MiniCheetah.self_collision_enabled = False
mini_cheetah.MiniCheetah.control_latency = 0.002
mini_cheetah.MiniCheetah.pd_latency = 0.0
mini_cheetah.MiniCheetah.motor_kp = [%ABDUCTION_P_GAIN, %HIP_P_GAIN, %KNEE_P_GAIN,
                            %ABDUCTION_P_GAIN, %HIP_P_GAIN, %KNEE_P_GAIN,
                            %ABDUCTION_P_GAIN, %HIP_P_GAIN, %KNEE_P_GAIN,
                            %ABDUCTION_P_GAIN, %HIP_P_GAIN, %KNEE_P_GAIN]
mini_cheetah.MiniCheetah.motor_kd = [%ABDUCTION_D_GAIN, %HIP_D_GAIN, %KNEE_D_GAIN,
                            %ABDUCTION_D_GAIN, %HIP_D_GAIN, %KNEE_D_GAIN,
                            %ABDUCTION_D_GAIN, %HIP_D_GAIN, %KNEE_D_GAIN,
                            %ABDUCTION_D_GAIN, %HIP_D_GAIN, %KNEE_D_GAIN]

locomotion_gym_env.LocomotionGymEnv.robot_class = @mini_cheetah.MiniCheetah

