import numpy as np
from gym.envs.mujoco import AntEnv  # AntV2
from gym.envs.mujoco.humanoid_v3 import HumanoidEnv



class TruncatedAnt(AntEnv):
    def _get_obs(self):
        external_force_zeros = np.zeros_like(
            np.clip(self.sim.data.cfrc_ext, -1, 1).flat
        )
        return np.concatenate(
            [self.sim.data.qpos.flat[2:], self.sim.data.qvel.flat, external_force_zeros]
        )


class TruncatedHumanoidV3(HumanoidEnv):
    def _get_obs(self):
        position = self.sim.data.qpos.flat.copy()
        velocity = self.sim.data.qvel.flat.copy()

        com_inertia = self.sim.data.cinert.flat.copy()
        com_velocity = self.sim.data.cvel.flat.copy()

        actuator_forces = self.sim.data.qfrc_actuator.flat.copy()
        external_contact_forces = self.sim.data.cfrc_ext.flat.copy()
        zero_external_forces = np.zeros_like(external_contact_forces)

        if self._exclude_current_positions_from_observation:
            position = position[2:]

        return np.concatenate(
            (
                position,
                velocity,
                com_inertia,
                com_velocity,
                actuator_forces,
                zero_external_forces,
            )
        )
