# Path: gym3/env.py
# class Env:
#     """
#     An interface for reinforcement learning environments.
# 
#     :param ob_space: ValType representing the valid observations generated by the environment
#     :param ac_space: ValType representing the valid actions that an agent can take in the environment
#     :param num: number of simultaneous episodes
#     """
# 
#     def __init__(self, ob_space: ValType, ac_space: ValType, num: int):
#         self.ob_space = ob_space
#         self.ac_space = ac_space
#         self.num = num
# 
#     def observe(self) -> Tuple[Any, Any, Any]:
#         """
#         Structured data that's accessible to the agent
# 
#         Can be called zero or more times per act() call; it is idempotent and does
#         not affect the environment.  Returns the observation after the last act()
#         call has affected the environment or the initial state if no act()
#         call has been made.
# 
#         The following snippet shows what the start of a rollout looks like, and
#         the usual convention for indexing the values:
# 
#         env = Env()
#         # initial reward is ignored; didn't depend on actions
#         _reward, ob_0, first_0 = env.observe()
#         env.act(ac_0)
#         reward_0, ob_1, first_1 = env.observe()
#         env.act(ac_1)
# 
#         Note that the initial reward `_reward` is ignored by algorithms, since it
#         didn't depend on the actions. And that first_0 is always True since the
#         environment was just created.
# 
#         :returns:
#             reward: (float array) last reward, shaped (num,)
#             ob: observation matching ob_space, with (num,) prepended to the shape of each leaf
#             first: (bool array) whether new episode just started, shaped (num,)
#         """
#         raise NotImplementedError
# 
#     def get_info(self) -> List[Dict]:
#         """
#         Return unstructured diagnostics that aren't accessible to the agent
#         Per-episode stats, rendered images, etc.
#         Corresponds to same timestep as the observation from observe().
# 
#         :returns: a list of dictionaries with length `self.num`
#         """
#         return [{} for _ in range(self.num)]
# 
#     def act(self, ac: Any) -> None:
#         """
#         Apply an action
# 
#         :param ac: action matching ac_space, with (num,) prepended to the shape of each leaf
#         """
#         raise NotImplementedError
# 
#     def callmethod(
#         self, method: str, *args: Sequence[Any], **kwargs: Sequence[Any]
#     ) -> List[Any]:
#         """
#         Call a method on the underlying python object that offers the Gym3 interface.
#         By default, this is the Env instance (self), but this can be overridden if you are
#         encapsulating another object or if your wrapper wishes to handle the method call.
# 
#         :param method: name of method to call on the base object
#         :param *args, **kwargs: the value of each argument must be a list with length equal to `self.num`
# 
#         :returns: A list of objects, with length equal to `self.num`
#         """
#         return getattr(self, method)(*args, **kwargs)
# 
# Path: gym3/internal/misc.py
# def intprod(xs):
# def allsame(xs):
# 
# Path: gym3/types_np.py
# def concat(xs: Sequence[Any], axis: int = 0) -> Any:
#     """
#     Concatenate the (leaf) arrays from xs
# 
#     :param xs: list of trees with the same shape, where the leaf values are numpy arrays
#     :param axis: axis to concatenate along
#     """
#     return multimap(lambda *xs: np.concatenate(xs, axis=axis), *xs)
# 
# def split(x: Any, sections: Sequence[int]) -> Sequence[Any]:
#     """
#     Split the (leaf) arrays from the tree x
# 
#     Examples:
# 
#         split([1,2,3,4], [1,2,3,4]) => [[1], [2], [3], [4]]
#         split([1,2,3,4], [1,3,4]) => [[1], [2, 3], [4]]
# 
#     :param x: a tree where the leaf values are numpy arrays
#     :param sections: list of indices to split at (not sizes of each split)
# 
#     :returns: list of trees with length `len(sections)` with the same shape as x
#             where each leaf is the corresponding section of the leaf in x
#     """
#     result = []
#     start = 0
#     for end in sections:
#         select_tree = multimap(lambda arr: arr[start:end], x)
#         start = end
#         result.append(select_tree)
#     return result


from typing import Any, Dict, List, Sequence, Tuple
from gym3.env import Env
from gym3.internal import misc
from gym3.types_np import concat, split
import numpy as np




def _chunk_seq(x: Sequence, sizes: Sequence[int]) -> Sequence[Any]:
    result = []
    assert len(x) == sum(sizes), f"x has incorrect length {len(x)} != {sum(sizes)}"
    start = 0
    for size in sizes:
        end = start + size
        result.append(x[start:end])
        start = end
    return result


class ConcatEnv(Env):
    """
    Concatenate multiple environments into a single environment.

    :param envs: list of environments to concatenate, must all have the same ac_space and ob_space
    """


# Path: gym3/env.py
# class Env:
#     """
#     An interface for reinforcement learning environments.
# 
#     :param ob_space: ValType representing the valid observations generated by the environment
#     :param ac_space: ValType representing the valid actions that an agent can take in the environment
#     :param num: number of simultaneous episodes
#     """
# 
#     def __init__(self, ob_space: ValType, ac_space: ValType, num: int):
#         self.ob_space = ob_space
#         self.ac_space = ac_space
#         self.num = num
# 
#     def observe(self) -> Tuple[Any, Any, Any]:
#         """
#         Structured data that's accessible to the agent
# 
#         Can be called zero or more times per act() call; it is idempotent and does
#         not affect the environment.  Returns the observation after the last act()
#         call has affected the environment or the initial state if no act()
#         call has been made.
# 
#         The following snippet shows what the start of a rollout looks like, and
#         the usual convention for indexing the values:
# 
#         env = Env()
#         # initial reward is ignored; didn't depend on actions
#         _reward, ob_0, first_0 = env.observe()
#         env.act(ac_0)
#         reward_0, ob_1, first_1 = env.observe()
#         env.act(ac_1)
# 
#         Note that the initial reward `_reward` is ignored by algorithms, since it
#         didn't depend on the actions. And that first_0 is always True since the
#         environment was just created.
# 
#         :returns:
#             reward: (float array) last reward, shaped (num,)
#             ob: observation matching ob_space, with (num,) prepended to the shape of each leaf
#             first: (bool array) whether new episode just started, shaped (num,)
#         """
#         raise NotImplementedError
# 
#     def get_info(self) -> List[Dict]:
#         """
#         Return unstructured diagnostics that aren't accessible to the agent
#         Per-episode stats, rendered images, etc.
#         Corresponds to same timestep as the observation from observe().
# 
#         :returns: a list of dictionaries with length `self.num`
#         """
#         return [{} for _ in range(self.num)]
# 
#     def act(self, ac: Any) -> None:
#         """
#         Apply an action
# 
#         :param ac: action matching ac_space, with (num,) prepended to the shape of each leaf
#         """
#         raise NotImplementedError
# 
#     def callmethod(
#         self, method: str, *args: Sequence[Any], **kwargs: Sequence[Any]
#     ) -> List[Any]:
#         """
#         Call a method on the underlying python object that offers the Gym3 interface.
#         By default, this is the Env instance (self), but this can be overridden if you are
#         encapsulating another object or if your wrapper wishes to handle the method call.
# 
#         :param method: name of method to call on the base object
#         :param *args, **kwargs: the value of each argument must be a list with length equal to `self.num`
# 
#         :returns: A list of objects, with length equal to `self.num`
#         """
#         return getattr(self, method)(*args, **kwargs)
# 
# Path: gym3/internal/misc.py
# def intprod(xs):
# def allsame(xs):
# 
# Path: gym3/types_np.py
# def concat(xs: Sequence[Any], axis: int = 0) -> Any:
#     """
#     Concatenate the (leaf) arrays from xs
# 
#     :param xs: list of trees with the same shape, where the leaf values are numpy arrays
#     :param axis: axis to concatenate along
#     """
#     return multimap(lambda *xs: np.concatenate(xs, axis=axis), *xs)
# 
# def split(x: Any, sections: Sequence[int]) -> Sequence[Any]:
#     """
#     Split the (leaf) arrays from the tree x
# 
#     Examples:
# 
#         split([1,2,3,4], [1,2,3,4]) => [[1], [2], [3], [4]]
#         split([1,2,3,4], [1,3,4]) => [[1], [2, 3], [4]]
# 
#     :param x: a tree where the leaf values are numpy arrays
#     :param sections: list of indices to split at (not sizes of each split)
# 
#     :returns: list of trees with length `len(sections)` with the same shape as x
#             where each leaf is the corresponding section of the leaf in x
#     """
#     result = []
#     start = 0
#     for end in sections:
#         select_tree = multimap(lambda arr: arr[start:end], x)
#         start = end
#         result.append(select_tree)
#     return result


# Path: gym3/concat.py
from typing import Any, Dict, List, Sequence, Tuple
from gym3.env import Env
from gym3.internal import misc
from gym3.types_np import concat, split
import numpy as np





def _chunk_seq(x: Sequence, sizes: Sequence[int]) -> Sequence[Any]:
    result = []
    assert len(x) == sum(sizes), f"x has incorrect length {len(x)} != {sum(sizes)}"
    start = 0
    for size in sizes:
        end = start + size
        result.append(x[start:end])
        start = end
    return result


class ConcatEnv(Env):
    """
    Concatenate multiple environments into a single environment.

    :param envs: list of environments to concatenate, must all have the same ac_space and ob_space
    """


    def __init__(self, envs: Sequence[Env]):