# Path: meltingpot/python/utils/substrates/wrappers/discrete_action_wrapper.py
# T = TypeVar('T')
# def _validate_action(
#     action: Mapping[str, np.ndarray],
#     action_spec: Mapping[str, dm_env.specs.Array]) -> None:
# def _validate_action_table(
#     action_table: Sequence[Mapping[str, np.ndarray]],
#     action_spec: Mapping[str, dm_env.specs.Array]) -> None:
# def _immutable_action(
#     action: Mapping[str, Numeric],
#     action_spec: Mapping[str, dm_env.specs.Array],
# ) -> Mapping[str, np.ndarray]:
# def _immutable_action_table(
#     action_table: Sequence[Mapping[str, Numeric]],
#     action_spec: Mapping[str, dm_env.specs.Array],
# ) -> Sequence[Mapping[str, np.ndarray]]:
#   def __init__(self, env, action_table: Sequence[Mapping[str, Numeric]]):
#   def step(self, action: Sequence[int]):
#   def action_spec(self) -> Sequence[dm_env.specs.DiscreteArray]:
# class Wrapper(base.Wrapper):


from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from meltingpot.python.utils.substrates.wrappers import discrete_action_wrapper
import dm_env
import numpy as np
import dmlab2d
#     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.
"""Tests for discrete_action_wrapper."""





MOVE_SPEC = dm_env.specs.BoundedArray(
    shape=(), minimum=0, maximum=3, dtype=np.int8, name='MOVE')
TURN_SPEC = dm_env.specs.BoundedArray(
    shape=(), minimum=0, maximum=3, dtype=np.int8, name='TURN')
VALID_VALUE_0 = np.zeros([], dtype=np.int8)
VALID_VALUE_1 = np.array(3, dtype=np.int8)
INVALID_VALUE = np.array(4, dtype=np.int8)


class Lab2DToListsWrapperTest(parameterized.TestCase):

  def test_valid_set(self):
    env = mock.Mock(spec_set=dmlab2d.Environment)
    env.action_spec.return_value = [
        {'MOVE': MOVE_SPEC, 'TURN': TURN_SPEC},
        {'MOVE': MOVE_SPEC, 'TURN': TURN_SPEC},
    ]

# Path: meltingpot/python/utils/substrates/wrappers/discrete_action_wrapper.py
# T = TypeVar('T')
# def _validate_action(
#     action: Mapping[str, np.ndarray],
#     action_spec: Mapping[str, dm_env.specs.Array]) -> None:
# def _validate_action_table(
#     action_table: Sequence[Mapping[str, np.ndarray]],
#     action_spec: Mapping[str, dm_env.specs.Array]) -> None:
# def _immutable_action(
#     action: Mapping[str, Numeric],
#     action_spec: Mapping[str, dm_env.specs.Array],
# ) -> Mapping[str, np.ndarray]:
# def _immutable_action_table(
#     action_table: Sequence[Mapping[str, Numeric]],
#     action_spec: Mapping[str, dm_env.specs.Array],
# ) -> Sequence[Mapping[str, np.ndarray]]:
#   def __init__(self, env, action_table: Sequence[Mapping[str, Numeric]]):
#   def step(self, action: Sequence[int]):
#   def action_spec(self) -> Sequence[dm_env.specs.DiscreteArray]:
# class Wrapper(base.Wrapper):


# Path: meltingpot/python/utils/substrates/wrappers/discrete_action_wrapper_test.py
from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from meltingpot.python.utils.substrates.wrappers import discrete_action_wrapper
import dm_env
import numpy as np
import dmlab2d

#     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.
"""Tests for discrete_action_wrapper."""





MOVE_SPEC = dm_env.specs.BoundedArray(
    shape=(), minimum=0, maximum=3, dtype=np.int8, name='MOVE')
TURN_SPEC = dm_env.specs.BoundedArray(
    shape=(), minimum=0, maximum=3, dtype=np.int8, name='TURN')
VALID_VALUE_0 = np.zeros([], dtype=np.int8)
VALID_VALUE_1 = np.array(3, dtype=np.int8)
INVALID_VALUE = np.array(4, dtype=np.int8)


class Lab2DToListsWrapperTest(parameterized.TestCase):

  def test_valid_set(self):
    env = mock.Mock(spec_set=dmlab2d.Environment)
    env.action_spec.return_value = [
        {'MOVE': MOVE_SPEC, 'TURN': TURN_SPEC},
        {'MOVE': MOVE_SPEC, 'TURN': TURN_SPEC},
    ]

    discrete_action_wrapper.Wrapper(env, action_table=[