# Path: android_env/components/action_type.py
# class ActionType(enum.IntEnum):
#   TOUCH = 0
#   LIFT = 1
#   REPEAT = 2
# 
# Path: android_env/wrappers/base_wrapper.py
# class BaseWrapper(dm_env.Environment):
#   def __init__(self, env):
#   def reset(self) -> dm_env.TimeStep:
#   def step(self, action: Any) -> dm_env.TimeStep:
#   def _reset_state(self):
#   def _process_action(self, action: Any) -> Any:
#   def _process_timestep(self, timestep: dm_env.TimeStep) -> dm_env.TimeStep:
#   def observation_spec(self) -> Dict[str, specs.Array]:
#   def action_spec(self) -> Dict[str, specs.Array]:
#   def _wrapper_logs(self) -> Dict[str, Any]:
#   def android_logs(self) -> Dict[str, Any]:
#   def raw_env(self):
#   def __getattr__(self, attr):
#   def close(self):


from typing import Optional, Sequence, Dict
from android_env.components import action_type
from android_env.wrappers import base_wrapper
from dm_env import specs
import dm_env
import numpy as np
    assert len(self._parent_action_spec) == 2
    assert not self._parent_action_spec['action_type'].shape
    assert self._parent_action_spec['touch_position'].shape == (2,)

  @property
  def num_actions(self) -> int:
    """Number of discrete actions."""

    if self._redundant_actions:
      return np.product(self._action_grid) * self._num_action_types
    else:
      return np.product(self._action_grid) + self._num_action_types - 1

  def step(self, action: Dict[str, int]) -> dm_env.TimeStep:
    """Take a step in the base environment."""

    return self._env.step(self._process_action(action))

  def _process_action(self, action: Dict[str, int]) -> Dict[str, np.ndarray]:
    """Transforms action so that it agrees with AndroidEnv's action spec."""

    return {
        'action_type':
            np.array(self._get_action_type(action['action_id']),
                     dtype=self._parent_action_spec['action_type'].dtype),
        'touch_position':
            np.array(self._get_touch_position(action['action_id']),
                     dtype=self._parent_action_spec['touch_position'].dtype)
    }


# Path: android_env/components/action_type.py
# class ActionType(enum.IntEnum):
#   TOUCH = 0
#   LIFT = 1
#   REPEAT = 2
# 
# Path: android_env/wrappers/base_wrapper.py
# class BaseWrapper(dm_env.Environment):
#   def __init__(self, env):
#   def reset(self) -> dm_env.TimeStep:
#   def step(self, action: Any) -> dm_env.TimeStep:
#   def _reset_state(self):
#   def _process_action(self, action: Any) -> Any:
#   def _process_timestep(self, timestep: dm_env.TimeStep) -> dm_env.TimeStep:
#   def observation_spec(self) -> Dict[str, specs.Array]:
#   def action_spec(self) -> Dict[str, specs.Array]:
#   def _wrapper_logs(self) -> Dict[str, Any]:
#   def android_logs(self) -> Dict[str, Any]:
#   def raw_env(self):
#   def __getattr__(self, attr):
#   def close(self):


# Path: android_env/wrappers/discrete_action_wrapper.py
from typing import Optional, Sequence, Dict
from android_env.components import action_type
from android_env.wrappers import base_wrapper
from dm_env import specs
import dm_env
import numpy as np

    assert len(self._parent_action_spec) == 2
    assert not self._parent_action_spec['action_type'].shape
    assert self._parent_action_spec['touch_position'].shape == (2,)

  @property
  def num_actions(self) -> int:
    """Number of discrete actions."""

    if self._redundant_actions:
      return np.product(self._action_grid) * self._num_action_types
    else:
      return np.product(self._action_grid) + self._num_action_types - 1

  def step(self, action: Dict[str, int]) -> dm_env.TimeStep:
    """Take a step in the base environment."""

    return self._env.step(self._process_action(action))

  def _process_action(self, action: Dict[str, int]) -> Dict[str, np.ndarray]:
    """Transforms action so that it agrees with AndroidEnv's action spec."""

    return {
        'action_type':
            np.array(self._get_action_type(action['action_id']),
                     dtype=self._parent_action_spec['action_type'].dtype),
        'touch_position':
            np.array(self._get_touch_position(action['action_id']),
                     dtype=self._parent_action_spec['touch_position'].dtype)
    }


  def _get_action_type(self, action_id: int) -> action_type.ActionType: