# Path: android_env/components/action_type.py
# class ActionType(enum.IntEnum):
#   TOUCH = 0
#   LIFT = 1
#   REPEAT = 2
# 
# Path: android_env/components/utils.py
# def touch_position_to_pixel_position(
#     touch_position: np.ndarray,
#     width_height: Sequence[int],
# ) -> Tuple[int, int]:
# def transpose_pixels(frame: np.ndarray) -> np.ndarray:
# def orient_pixels(
#     frame: np.ndarray,
#     orientation: adb_pb2.AdbCall.Rotate.Orientation) -> np.ndarray:
# 
# 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 Dict
from android_env.components import action_type
from android_env.components import utils
from android_env.wrappers import base_wrapper
from dm_env import specs
import dm_env
import numpy as np
    super().__init__(env)
    self._concat_to_pixels = concat_to_pixels
    self._screen_dimensions = self._env.observation_spec()['pixels'].shape[:2]

  def _process_timestep(self, timestep: dm_env.TimeStep) -> dm_env.TimeStep:
    observation = timestep.observation.copy()
    processed_observation = self._process_observation(observation)
    return timestep._replace(observation=processed_observation)

  def _process_observation(
      self, observation: Dict[str, np.ndarray]
  ) -> Dict[str, np.ndarray]:
    """Extends observation with last_action data."""
    processed_observation = observation.copy()
    last_action_layer = self._get_last_action_layer(observation['pixels'])
    if self._concat_to_pixels:
      pixels = observation['pixels'].copy()
      processed_pixels = np.dstack((pixels, last_action_layer))
      processed_observation['pixels'] = processed_pixels
    else:
      processed_observation['last_action'] = last_action_layer
    return processed_observation

  def _get_last_action_layer(self, pixels: np.ndarray) -> np.ndarray:
    """Makes sure the rescaling doesn't distort the last_action layer."""

    last_action = self._env.raw_action
    last_action_layer = np.zeros(self._screen_dimensions, dtype=pixels.dtype)

    if ('action_type' in last_action and

# Path: android_env/components/action_type.py
# class ActionType(enum.IntEnum):
#   TOUCH = 0
#   LIFT = 1
#   REPEAT = 2
# 
# Path: android_env/components/utils.py
# def touch_position_to_pixel_position(
#     touch_position: np.ndarray,
#     width_height: Sequence[int],
# ) -> Tuple[int, int]:
# def transpose_pixels(frame: np.ndarray) -> np.ndarray:
# def orient_pixels(
#     frame: np.ndarray,
#     orientation: adb_pb2.AdbCall.Rotate.Orientation) -> np.ndarray:
# 
# 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/last_action_wrapper.py
from typing import Dict
from android_env.components import action_type
from android_env.components import utils
from android_env.wrappers import base_wrapper
from dm_env import specs
import dm_env
import numpy as np

    super().__init__(env)
    self._concat_to_pixels = concat_to_pixels
    self._screen_dimensions = self._env.observation_spec()['pixels'].shape[:2]

  def _process_timestep(self, timestep: dm_env.TimeStep) -> dm_env.TimeStep:
    observation = timestep.observation.copy()
    processed_observation = self._process_observation(observation)
    return timestep._replace(observation=processed_observation)

  def _process_observation(
      self, observation: Dict[str, np.ndarray]
  ) -> Dict[str, np.ndarray]:
    """Extends observation with last_action data."""
    processed_observation = observation.copy()
    last_action_layer = self._get_last_action_layer(observation['pixels'])
    if self._concat_to_pixels:
      pixels = observation['pixels'].copy()
      processed_pixels = np.dstack((pixels, last_action_layer))
      processed_observation['pixels'] = processed_pixels
    else:
      processed_observation['last_action'] = last_action_layer
    return processed_observation

  def _get_last_action_layer(self, pixels: np.ndarray) -> np.ndarray:
    """Makes sure the rescaling doesn't distort the last_action layer."""

    last_action = self._env.raw_action
    last_action_layer = np.zeros(self._screen_dimensions, dtype=pixels.dtype)

    if ('action_type' in last_action and

        last_action['action_type'] == action_type.ActionType.TOUCH):