# -*- coding: utf-8 -*-
"""Go1 configuration for the RoboGauge MuJoCo evaluator."""

from robogauge.tasks.robots.go2.go2_config import Go2Config


class Go1Config(Go2Config):
    robot_name = 'go1'
    robot_class = 'Go1'

    class assets(Go2Config.assets):
        robot_xml = "{ROBOGAUGE_ROOT_DIR}/resources/robots/go1/go1.xml"
        foot_geom_names = ['FR', 'FL', 'RR', 'RL']

    class control(Go2Config.control):
        model_path = "{ROBOGAUGE_ROOT_DIR}/resources/models/go1/policy.pt"
        control_dt = 0.02
        p_gains = [28.0] * 12
        d_gains = [0.7] * 12

        # Both this MuJoCo model and the trained policy use FR, FL, RR, RL.
        # Go1Robot permuted Isaac Gym's FL, FR, RL, RR asset state into this
        # order during training, so applying that permutation again here
        # would swap every left/right leg pair.
        default_dof_pos = [
            -0.1, 0.8, -1.5,
             0.1, 0.8, -1.5,
            -0.1, 1.0, -1.5,
             0.1, 1.0, -1.5,
        ]
        mj2model_dof_indices = list(range(12))


class Go1TerrainConfig(Go1Config):
    class commands(Go1Config.commands):
        lin_vel_x = [-1.0, 1.0]
        lin_vel_y = [-1.0, 1.0]
        ang_vel_yaw = [-1.5, 1.5]