from legged_gym.envs.go2.go2_env import Go2Robot
import torch


class Go1Robot(Go2Robot):
    """Go1 asset with an explicit policy-to-asset joint permutation.

    Isaac Gym loads this asset as [FL, FR, RL, RR]. The permutation below
    presents observations to the policy as [FR, FL, RR, RL] and maps policy
    actions back to the Isaac Gym asset order. Existing Go1 checkpoints depend
    on this policy order.
    """

    # Policy order [FR, FL, RR, RL] -> asset order [FL, FR, RL, RR].
    _POLICY_TO_ASSET = [3, 4, 5, 0, 1, 2, 9, 10, 11, 6, 7, 8]

    def step(self, actions):
        asset_actions = torch.zeros_like(actions)
        asset_actions[:, self._POLICY_TO_ASSET] = actions
        return super().step(asset_actions)

    def compute_observations(self):
        super().compute_observations()

        # Go2 policy observation: angular velocity, gravity, commands,
        # joint-position error, joint velocity, previous action.
        for start in (9, 21, 33):
            self.obs_buf[:, start:start + 12] = self.obs_buf[:, start:start + 12][:, self._POLICY_TO_ASSET]

        # Keep the privileged teacher input consistent as well. Its first
        # three fields are base linear velocity, so joint fields start at 12.
        if self.privileged_obs_buf is not None:
            for start in (12, 24, 36, 52, 64):
                self.privileged_obs_buf[:, start:start + 12] = self.privileged_obs_buf[:, start:start + 12][:, self._POLICY_TO_ASSET]