# Path: meltingpot/python/utils/scenarios/wrappers/base.py
# class Wrapper(substrate.Substrate):
#   def __init__(self, env: substrate.Substrate):
#   def reset(self, *args, **kwargs) -> ...:
#   def step(self, *args, **kwargs) -> ...:
#   def reward_spec(self, *args, **kwargs) -> ...:
#   def discount_spec(self, *args, **kwargs) -> ...:
#   def observation_spec(self, *args, **kwargs) -> ...:
#   def action_spec(self, *args, **kwargs) -> ...:
#   def close(self, *args, **kwargs) -> ...:
#   def observation(self, *args, **kwargs) -> ...:
#   def events(self, *args, **kwargs) -> ...:
#   def list_property(self, *args, **kwargs) -> ...:
#   def write_property(self, *args, **kwargs) -> ...:
#   def read_property(self, *args, **kwargs) -> ...:
#   def observables(self, *args, **kwargs) -> ...:


import dm_env
import numpy as np
from meltingpot.python.utils.scenarios.wrappers import base
# Copyright 2020 DeepMind Technologies Limited.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Wrap environment, adding the agent's player slot to the observation."""



AGENT_SLOT = 'agent_slot'


def _augment_timestep(timestep: dm_env.TimeStep) -> dm_env.TimeStep:
  """Returns a new timestep with player slot added as a one-hot."""
  player_slots = np.eye(len(timestep.observation), dtype=np.float32)
  observations = [
      dict(observation, **{AGENT_SLOT: player_slots[n]})
      for n, observation in enumerate(timestep.observation)
  ]
  return timestep._replace(observation=observations)



# Path: meltingpot/python/utils/scenarios/wrappers/base.py
# class Wrapper(substrate.Substrate):
#   def __init__(self, env: substrate.Substrate):
#   def reset(self, *args, **kwargs) -> ...:
#   def step(self, *args, **kwargs) -> ...:
#   def reward_spec(self, *args, **kwargs) -> ...:
#   def discount_spec(self, *args, **kwargs) -> ...:
#   def observation_spec(self, *args, **kwargs) -> ...:
#   def action_spec(self, *args, **kwargs) -> ...:
#   def close(self, *args, **kwargs) -> ...:
#   def observation(self, *args, **kwargs) -> ...:
#   def events(self, *args, **kwargs) -> ...:
#   def list_property(self, *args, **kwargs) -> ...:
#   def write_property(self, *args, **kwargs) -> ...:
#   def read_property(self, *args, **kwargs) -> ...:
#   def observables(self, *args, **kwargs) -> ...:


# Path: meltingpot/python/utils/scenarios/wrappers/agent_slot_wrapper.py
import dm_env
import numpy as np
from meltingpot.python.utils.scenarios.wrappers import base

# Copyright 2020 DeepMind Technologies Limited.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Wrap environment, adding the agent's player slot to the observation."""



AGENT_SLOT = 'agent_slot'


def _augment_timestep(timestep: dm_env.TimeStep) -> dm_env.TimeStep:
  """Returns a new timestep with player slot added as a one-hot."""
  player_slots = np.eye(len(timestep.observation), dtype=np.float32)
  observations = [
      dict(observation, **{AGENT_SLOT: player_slots[n]})
      for n, observation in enumerate(timestep.observation)
  ]
  return timestep._replace(observation=observations)



class Wrapper(base.Wrapper):