import collections


def flatten_observation_dim(observation_dim, observation_excluded=()):
  """Flattens the observation dimensions to an array.

  If observation_excluded is passed in, it will still return a dictionary,
  which includes all the (key, observation_dict[key]) in observation_excluded,
  and ('other': the flattened array).

  Args:
    observation_dim: A dictionary of all the observation dimensions.
    observation_excluded: A list/tuple of all the keys of the observations to be
      ignored during flattening.

  Returns:
    An array or a dictionary of observation dimensions based on whether
      observation_excluded is empty.
  """
  if not isinstance(observation_excluded, (list, tuple)):
    observation_excluded = [observation_excluded]
  observation_dims = 0
  for key, value in observation_dim.items():
    if key not in observation_excluded:
      observation_dims += value
  if not observation_excluded:
    return observation_dims
  else:
    dim_dict_after_flatten = {"other": observation_dims}
    for key in observation_excluded:
      if key in observation_dim:
        dim_dict_after_flatten[key] = observation_dim[key]
    return collections.OrderedDict(sorted(list(dim_dict_after_flatten.items())))