# Path: android_env/components/action_type.py
# class ActionType(enum.IntEnum):
#   TOUCH = 0
#   LIFT = 1
#   REPEAT = 2


from typing import Dict
from android_env.components import action_type
from android_env.proto import task_pb2
from dm_env import specs
import dm_env
import numpy as np
    task_pb2.ArraySpec.DataType.UINT8: np.uint8,
    task_pb2.ArraySpec.DataType.UINT16: np.uint16,
    task_pb2.ArraySpec.DataType.UINT32: np.uint32,
    task_pb2.ArraySpec.DataType.UINT64: np.uint64,
    task_pb2.ArraySpec.DataType.BOOL: np.bool_,
    task_pb2.ArraySpec.DataType.STRING_U1: np.dtype(('U1')),
    task_pb2.ArraySpec.DataType.STRING_U16: np.dtype(('<U16')),
    task_pb2.ArraySpec.DataType.STRING_U25: np.dtype(('<U25')),
    task_pb2.ArraySpec.DataType.STRING_U250: np.dtype(('<U250')),
}


def base_action_spec(num_fingers: int = 1) -> Dict[str, specs.Array]:
  """Default action spec for AndroidEnv.

  Args:
    num_fingers: Number of virtual fingers of the agent.
  Returns:
    A dict of action specs, each item corresponding to a virtual finger.
    action_type: An integer of type ActionType: TOUCH=0, LIFT=1, REPEAT=2
    touch_position: Position [x, y] of the touch action, where x, y are float
      values between 0.0 and 1.0 corresponding to the relative position on the
      screen. IGNORED when (action_type != ActionType.TOUCH).
    action_type_i: Action type for additional fingers (i>1).
    touch_position_i: Touch position for additional fingers (i>1).
  """

  action_spec = {
      'action_type':
          specs.DiscreteArray(

# Path: android_env/components/action_type.py
# class ActionType(enum.IntEnum):
#   TOUCH = 0
#   LIFT = 1
#   REPEAT = 2


# Path: android_env/components/specs.py
from typing import Dict
from android_env.components import action_type
from android_env.proto import task_pb2
from dm_env import specs
import dm_env
import numpy as np

    task_pb2.ArraySpec.DataType.UINT8: np.uint8,
    task_pb2.ArraySpec.DataType.UINT16: np.uint16,
    task_pb2.ArraySpec.DataType.UINT32: np.uint32,
    task_pb2.ArraySpec.DataType.UINT64: np.uint64,
    task_pb2.ArraySpec.DataType.BOOL: np.bool_,
    task_pb2.ArraySpec.DataType.STRING_U1: np.dtype(('U1')),
    task_pb2.ArraySpec.DataType.STRING_U16: np.dtype(('<U16')),
    task_pb2.ArraySpec.DataType.STRING_U25: np.dtype(('<U25')),
    task_pb2.ArraySpec.DataType.STRING_U250: np.dtype(('<U250')),
}


def base_action_spec(num_fingers: int = 1) -> Dict[str, specs.Array]:
  """Default action spec for AndroidEnv.

  Args:
    num_fingers: Number of virtual fingers of the agent.
  Returns:
    A dict of action specs, each item corresponding to a virtual finger.
    action_type: An integer of type ActionType: TOUCH=0, LIFT=1, REPEAT=2
    touch_position: Position [x, y] of the touch action, where x, y are float
      values between 0.0 and 1.0 corresponding to the relative position on the
      screen. IGNORED when (action_type != ActionType.TOUCH).
    action_type_i: Action type for additional fingers (i>1).
    touch_position_i: Touch position for additional fingers (i>1).
  """

  action_spec = {
      'action_type':
          specs.DiscreteArray(

              num_values=len(action_type.ActionType),