from legged_gym.envs.go2.go2_config import GO2Cfg, GO2CfgMoECTS


class GO1Cfg(GO2Cfg):
    class init_state(GO2Cfg.init_state):
        pos = [0.0, 0.0, 0.34]

    class control(GO2Cfg.control):
        stiffness = {'joint': 28.0}
        damping = {'joint': 0.7}
        action_scale = 0.25
        decimation = 4

    class asset(GO2Cfg.asset):
        file = '{LEGGED_GYM_ROOT_DIR}/resources/robots/go1/urdf/go1.urdf'
        name = 'go1'
        self_collisions = 1

    class rewards(GO2Cfg.rewards):
        base_height_target = 0.30


class GO1CfgMoECTS(GO2CfgMoECTS):
    class runner(GO2CfgMoECTS.runner):
        experiment_name = 'go1_moe_cts'
        run_name = 'from_scratch_5000'
        max_iterations = 5000
        save_interval = 500
        resume = False
        load_run = -1
        checkpoint = -1
        resume_path = None


class GO1StairsCfg(GO1Cfg):
    class terrain(GO1Cfg.terrain):
        # wave, slope, rough slope, stairs up, stairs down, obstacles,
        # stepping stones, gap, flat
        terrain_proportions = [0.03, 0.06, 0.03, 0.50, 0.20, 0.08, 0.0, 0.0, 0.10]

    class commands(GO1Cfg.commands):
        # The 10k checkpoint has only seen +/-0.5 m/s because the original
        # curriculum does not widen until iteration 20000. Cover the 0.8 m/s
        # RoboGauge stair command immediately during this fine-tune.
        command_range_curriculum = []

        class ranges(GO1Cfg.commands.ranges):
            lin_vel_x = [-1.0, 1.0]
            lin_vel_y = [-0.5, 0.5]
            ang_vel_yaw = [-1.0, 1.0]


class GO1StairsCfgMoECTS(GO1CfgMoECTS):
    class runner(GO1CfgMoECTS.runner):
        experiment_name = 'go1_moe_cts'
        run_name = 'stairs_finetune_10000_to_15000'
        max_iterations = 5000
        save_interval = 500