# Path: meltingpot/python/utils/scenarios/wrappers/agent_slot_wrapper.py
# AGENT_SLOT = 'agent_slot'
# def _augment_timestep(timestep: dm_env.TimeStep) -> dm_env.TimeStep:
#   def reset(self):
#   def step(self, actions):
#   def observation_spec(self):
# class Wrapper(base.Wrapper):
# 
# 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) -> ...:


from unittest import mock
from absl.testing import absltest
from meltingpot.python.utils.scenarios.wrappers import agent_slot_wrapper
from meltingpot.python.utils.scenarios.wrappers import base
import dm_env
import numpy as np
RGB_SPEC = dm_env.specs.Array(shape=(8, 8, 3), dtype=np.int8)
RGB_VALUE = np.ones((8, 8, 3), np.int8)


class AgentSlotWrapperTest(absltest.TestCase):

  def test_augment_timestep(self):
    timestep = dm_env.restart([{'RGB': RGB_VALUE}] * 3)
    actual = agent_slot_wrapper._augment_timestep(timestep)
    expected = dm_env.restart([
        {
            'RGB': RGB_VALUE,
            AGENT_SLOT: np.array([1, 0, 0], dtype=np.float32)
        },
        {
            'RGB': RGB_VALUE,
            AGENT_SLOT: np.array([0, 1, 0], dtype=np.float32)
        },
        {
            'RGB': RGB_VALUE,
            AGENT_SLOT: np.array([0, 0, 1], dtype=np.float32)
        },
    ])

    with self.subTest('not_inplace'):
      self.assertIsNot(actual.observation, timestep.observation)
    with self.subTest('not_inplace'):
      np.testing.assert_equal(actual, expected)

  def test_added_slot(self):

# Path: meltingpot/python/utils/scenarios/wrappers/agent_slot_wrapper.py
# AGENT_SLOT = 'agent_slot'
# def _augment_timestep(timestep: dm_env.TimeStep) -> dm_env.TimeStep:
#   def reset(self):
#   def step(self, actions):
#   def observation_spec(self):
# class Wrapper(base.Wrapper):
# 
# 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_test.py
from unittest import mock
from absl.testing import absltest
from meltingpot.python.utils.scenarios.wrappers import agent_slot_wrapper
from meltingpot.python.utils.scenarios.wrappers import base
import dm_env
import numpy as np

RGB_SPEC = dm_env.specs.Array(shape=(8, 8, 3), dtype=np.int8)
RGB_VALUE = np.ones((8, 8, 3), np.int8)


class AgentSlotWrapperTest(absltest.TestCase):

  def test_augment_timestep(self):
    timestep = dm_env.restart([{'RGB': RGB_VALUE}] * 3)
    actual = agent_slot_wrapper._augment_timestep(timestep)
    expected = dm_env.restart([
        {
            'RGB': RGB_VALUE,
            AGENT_SLOT: np.array([1, 0, 0], dtype=np.float32)
        },
        {
            'RGB': RGB_VALUE,
            AGENT_SLOT: np.array([0, 1, 0], dtype=np.float32)
        },
        {
            'RGB': RGB_VALUE,
            AGENT_SLOT: np.array([0, 0, 1], dtype=np.float32)
        },
    ])

    with self.subTest('not_inplace'):
      self.assertIsNot(actual.observation, timestep.observation)
    with self.subTest('not_inplace'):
      np.testing.assert_equal(actual, expected)

  def test_added_slot(self):

    env = mock.Mock(spec_set=base.Substrate)