# 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 Union, Dict, Any
from android_env.wrappers import base_wrapper
from dm_env import specs
import dm_env
import numpy as np

"""Wraps the AndroidEnv environment to make its interface flat."""



RGB_CHANNELS = (0, 1, 2)


def _extract_screen_pixels(obs: np.ndarray):
  """Get only screen pixels by removing previous action layer."""
  is_grayscale_image = obs.shape[-1] == 2
  if is_grayscale_image:
    return np.expand_dims(obs[..., 0], -1)
  return obs[..., RGB_CHANNELS]


def _get_no_action_observation_spec(obs_spec: specs.BoundedArray):
  """Create an observation spec without the action layer."""
  shape = np.array(obs_spec.shape)
  shape[2] -= 1
  minimum = obs_spec.minimum
  maximum = obs_spec.maximum
  is_scalar = lambda x: np.isscalar(x) or np.ndim(x) == 0
  if not is_scalar(minimum):
    minimum = _extract_screen_pixels(minimum)
  if not is_scalar(maximum):
    maximum = _extract_screen_pixels(maximum)
  return obs_spec.replace(shape=shape, minimum=minimum, maximum=maximum)



# 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/flat_interface_wrapper.py
from typing import Union, Dict, Any
from android_env.wrappers import base_wrapper
from dm_env import specs
import dm_env
import numpy as np


"""Wraps the AndroidEnv environment to make its interface flat."""



RGB_CHANNELS = (0, 1, 2)


def _extract_screen_pixels(obs: np.ndarray):
  """Get only screen pixels by removing previous action layer."""
  is_grayscale_image = obs.shape[-1] == 2
  if is_grayscale_image:
    return np.expand_dims(obs[..., 0], -1)
  return obs[..., RGB_CHANNELS]


def _get_no_action_observation_spec(obs_spec: specs.BoundedArray):
  """Create an observation spec without the action layer."""
  shape = np.array(obs_spec.shape)
  shape[2] -= 1
  minimum = obs_spec.minimum
  maximum = obs_spec.maximum
  is_scalar = lambda x: np.isscalar(x) or np.ndim(x) == 0
  if not is_scalar(minimum):
    minimum = _extract_screen_pixels(minimum)
  if not is_scalar(maximum):
    maximum = _extract_screen_pixels(maximum)
  return obs_spec.replace(shape=shape, minimum=minimum, maximum=maximum)



class FlatInterfaceWrapper(base_wrapper.BaseWrapper):