import pybullet_envs.minitaur.envs_v2.locomotion_gym_config
import pybullet_envs.minitaur.envs_v2.sensors.robot_sensors

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()

robot_sensors.IMUSensor.channels = ["R", "P", "dR", "dP"]
robot_sensors.IMUSensor.noisy_reading = %NOISY_READING
robot_sensors.IMUSensor.lower_bound = [-6.28318548203, -6.28318548203,
                                       -6283.18554688, -6283.18554688]
robot_sensors.IMUSensor.upper_bound = [6.28318548203, 6.28318548203,
                                       6283.18554688, 6283.18554688]

robot_sensors.MotorAngleSensor.num_motors = %NUM_MOTORS
robot_sensors.MotorAngleSensor.noisy_reading = %NOISY_READING
robot_sensors.MotorAngleSensor.lower_bound = -6.28318548203
robot_sensors.MotorAngleSensor.upper_bound = 6.28318548203

sensors = [@robot_sensors.IMUSensor(), @robot_sensors.MotorAngleSensor()]

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()]
