from typing import Union
import torch
from typing import Sequence
from typing import Mapping
from typing import List
from typing import Any


def default_decollate(
        batch: Union[torch.Tensor, Sequence, Mapping],
        ignore: List[str] = ['prev_state', 'prev_actor_state', 'prev_critic_state']
) -> List[Any]:
    """
    Overview:
        Drag out batch_size collated data's batch size to decollate it,
        which is the reverse operation of ``default_collate``.
    Arguments:
        - batch (:obj:`Union[torch.Tensor, Sequence, Mapping]`): can refer to the Returns of ``default_collate``
        - ignore(:obj:`List[str]`): a list of names to be ignored, only function if input ``batch`` is a dict. \
            If key is in this list, its value would stay the same with no decollation.
    Returns:
        - ret (:obj:`List[Any]`): a list with B elements.
    """
    if isinstance(batch, torch.Tensor):
        batch = torch.split(batch, 1, dim=0)
        # squeeze if original batch's shape is like (B, dim1, dim2, ...);
        # otherwise directly return the list.
        if len(batch[0].shape) > 1:
            batch = [elem.squeeze(0) for elem in batch]
        return list(batch)
    elif isinstance(batch, Sequence):
        return list(zip(*[default_decollate(e) for e in batch]))
    elif isinstance(batch, Mapping):
        tmp = {k: v if k in ignore else default_decollate(v) for k, v in batch.items()}
        B = len(list(tmp.values())[0])
        return [{k: tmp[k][i] for k in tmp.keys()} for i in range(B)]

    raise TypeError("not support batch type: {}".format(type(batch)))