# Path: gym3/types.py
# class Discrete(ScalarType):
#     """
#     A scalar that can represent discrete (integer) values
# 
#     :param n: number of discrete values that this scalar can assume, which are required to be in the range [0, n)
#     :param bit_width: number of bits for the integer dtype that will be used to store values
#     :param signed: Set to False to use an unsigned integer type
#     """
# 
#     def __init__(self, n: int, dtype_name: str = "int64") -> None:
#         self.n = int(n)
#         self.dtype_name = dtype_name
#         assert self.dtype_name in INTEGER_DTYPE_NAMES
#         max_value = DTYPE_NAME_TO_MAX_VALUE[self.dtype_name]
#         assert (
#             n <= max_value + 1
#         ), f"{n} cannot be greater than {max_value + 1} for dtype_name={dtype_name}"
# 
#     def __repr__(self):
#         return f"<Discrete: {self}>"
# 
#     def __str__(self):
#         return f"D{self.n}"
# 
# class Real(ScalarType):
#     """
#     A scalar that can represent continuous values
#     """
# 
#     def __init__(self, dtype_name: str = "float32"):
#         self.dtype_name = dtype_name
# 
#     def __repr__(self):
#         return "<Real>"
# 
#     def __str__(self):
#         return "R"
# 
# class TensorType(ValType):
#     """
#     A tensor value type
# 
#     :param eltype: instance of ScalarType subclass that represents the types of values in this tensor
#     :param shape: shape of the tensor as a tuple of ints
#     """
# 
#     def __init__(self, eltype: ScalarType, shape: Tuple) -> None:
#         assert isinstance(shape, tuple)
#         self.eltype = eltype
#         self.shape = shape
# 
#     @property
#     def ndim(self):
#         """
#         Number of dimensions of the tensor
#         """
#         return len(self.shape)
# 
#     @property
#     def size(self):
#         """
#         Number of elements of the tensor
#         """
#         return int(misc.intprod(self.shape))
# 
#     def __repr__(self):
#         return f"<TensorType: {self}>"
# 
#     def __str__(self):
#         shape_str = ",".join(map(str, self.shape))
#         return f"{self.eltype}[{shape_str}]"
# 
# class ValType:
#     """
#     Tensor or combination of tensors
#     """
# 
#     __eq__ = pod_equals
# 
# def multimap(f: Callable, *xs: Any) -> Any:
#     """
#     Apply f at each leaf of the list of trees
# 
#     A tree is:
#         * a (possibly nested) dict
#         * a (possibly nested) DictType
#         * any other object (a leaf value)
# 
#     `{"a": 1}`, `{"a": {"b": 2}}`, and `3` are all valid trees, where the leaf values
#     are the integers
# 
#     :param f: function to call at each leaf, must take len(xs) arguments
#     :param xs: a list of trees, all with the same structure
# 
#     :returns: A tree of the same structure, where each leaf contains f's return value.
#     """
#     first = xs[0]
#     if isinstance(first, dict) or isinstance(first, DictType):
#         assert all(isinstance(x, dict) or isinstance(x, DictType) for x in xs)
#         assert all(sorted(x.keys()) == sorted(first.keys()) for x in xs)
#         return {k: multimap(f, *(x[k] for x in xs)) for k in sorted(first.keys())}
#     else:
#         return f(*xs)


from functools import partial
from typing import Any, Sequence, Tuple
from gym3.types import Discrete, Real, TensorType, ValType, multimap
import torch as th

def stack(xs: Sequence[Any], dim: int = 0) -> Any:
    """
    Stack the (leaf) tensors from xs

    :param xs: list of trees with the same shape, where the leaf values are torch tensors
    :param dim: dimension to stack along
    """
    return multimap(lambda *xs: th.stack(xs, dim=dim), *xs)


def split(x: Any, sections: Sequence[int]) -> Any:
    """
    Split the (leaf) tensors 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 torch tensors
    :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
    """
    # split each leaf and select the correct component
    result = []
    start = 0
    for end in sections:

# Path: gym3/types.py
# class Discrete(ScalarType):
#     """
#     A scalar that can represent discrete (integer) values
# 
#     :param n: number of discrete values that this scalar can assume, which are required to be in the range [0, n)
#     :param bit_width: number of bits for the integer dtype that will be used to store values
#     :param signed: Set to False to use an unsigned integer type
#     """
# 
#     def __init__(self, n: int, dtype_name: str = "int64") -> None:
#         self.n = int(n)
#         self.dtype_name = dtype_name
#         assert self.dtype_name in INTEGER_DTYPE_NAMES
#         max_value = DTYPE_NAME_TO_MAX_VALUE[self.dtype_name]
#         assert (
#             n <= max_value + 1
#         ), f"{n} cannot be greater than {max_value + 1} for dtype_name={dtype_name}"
# 
#     def __repr__(self):
#         return f"<Discrete: {self}>"
# 
#     def __str__(self):
#         return f"D{self.n}"
# 
# class Real(ScalarType):
#     """
#     A scalar that can represent continuous values
#     """
# 
#     def __init__(self, dtype_name: str = "float32"):
#         self.dtype_name = dtype_name
# 
#     def __repr__(self):
#         return "<Real>"
# 
#     def __str__(self):
#         return "R"
# 
# class TensorType(ValType):
#     """
#     A tensor value type
# 
#     :param eltype: instance of ScalarType subclass that represents the types of values in this tensor
#     :param shape: shape of the tensor as a tuple of ints
#     """
# 
#     def __init__(self, eltype: ScalarType, shape: Tuple) -> None:
#         assert isinstance(shape, tuple)
#         self.eltype = eltype
#         self.shape = shape
# 
#     @property
#     def ndim(self):
#         """
#         Number of dimensions of the tensor
#         """
#         return len(self.shape)
# 
#     @property
#     def size(self):
#         """
#         Number of elements of the tensor
#         """
#         return int(misc.intprod(self.shape))
# 
#     def __repr__(self):
#         return f"<TensorType: {self}>"
# 
#     def __str__(self):
#         shape_str = ",".join(map(str, self.shape))
#         return f"{self.eltype}[{shape_str}]"
# 
# class ValType:
#     """
#     Tensor or combination of tensors
#     """
# 
#     __eq__ = pod_equals
# 
# def multimap(f: Callable, *xs: Any) -> Any:
#     """
#     Apply f at each leaf of the list of trees
# 
#     A tree is:
#         * a (possibly nested) dict
#         * a (possibly nested) DictType
#         * any other object (a leaf value)
# 
#     `{"a": 1}`, `{"a": {"b": 2}}`, and `3` are all valid trees, where the leaf values
#     are the integers
# 
#     :param f: function to call at each leaf, must take len(xs) arguments
#     :param xs: a list of trees, all with the same structure
# 
#     :returns: A tree of the same structure, where each leaf contains f's return value.
#     """
#     first = xs[0]
#     if isinstance(first, dict) or isinstance(first, DictType):
#         assert all(isinstance(x, dict) or isinstance(x, DictType) for x in xs)
#         assert all(sorted(x.keys()) == sorted(first.keys()) for x in xs)
#         return {k: multimap(f, *(x[k] for x in xs)) for k in sorted(first.keys())}
#     else:
#         return f(*xs)


# Path: gym3/types_th.py
from functools import partial
from typing import Any, Sequence, Tuple
from gym3.types import Discrete, Real, TensorType, ValType, multimap
import torch as th


def stack(xs: Sequence[Any], dim: int = 0) -> Any:
    """
    Stack the (leaf) tensors from xs

    :param xs: list of trees with the same shape, where the leaf values are torch tensors
    :param dim: dimension to stack along
    """
    return multimap(lambda *xs: th.stack(xs, dim=dim), *xs)


def split(x: Any, sections: Sequence[int]) -> Any:
    """
    Split the (leaf) tensors 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 torch tensors
    :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
    """
    # split each leaf and select the correct component
    result = []
    start = 0
    for end in sections:

        select_tree = multimap(lambda t: t[start:end], x)