from collections.abc import Sequence, Mapping
from typing import List, Dict, Union, Any

import torch
import treetensor.torch as ttorch
import re
import collections.abc as container_abcs
from ding.compatibility import torch_ge_131

int_classes = int
string_classes = (str, bytes)
np_str_obj_array_pattern = re.compile(r'[SaUO]')

default_collate_err_msg_format = (
    "default_collate: batch must contain tensors, numpy arrays, numbers, "
    "dicts or lists; found {}"
)


def ttorch_collate(x, json: bool = False, cat_1dim: bool = True):
    """
    Overview:
        Collates a list of tensors or nested dictionaries of tensors into a single tensor or nested \
            dictionary of tensors.

    Arguments:
        - x : The input list of tensors or nested dictionaries of tensors.
        - json (:obj:`bool`): If True, converts the output to JSON format. Defaults to False.
        - cat_1dim (:obj:`bool`): If True, concatenates tensors with shape (B, 1) along the last dimension. \
            Defaults to True.

    Returns:
        The collated output tensor or nested dictionary of tensors.

    Examples:
        >>> # case 1: Collate a list of tensors
        >>> tensors = [torch.tensor([1, 2, 3]), torch.tensor([4, 5, 6]), torch.tensor([7, 8, 9])]
        >>> collated = ttorch_collate(tensors)
        collated = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
        >>> # case 2: Collate a nested dictionary of tensors
        >>> nested_dict = {
                'a': torch.tensor([1, 2, 3]),
                'b': torch.tensor([4, 5, 6]),
                'c': torch.tensor([7, 8, 9])
            }
        >>> collated = ttorch_collate(nested_dict)
        collated = {
            'a': torch.tensor([1, 2, 3]),
            'b': torch.tensor([4, 5, 6]),
            'c': torch.tensor([7, 8, 9])
        }
        >>> # case 3: Collate a list of nested dictionaries of tensors
        >>> nested_dicts = [
                {'a': torch.tensor([1, 2, 3]), 'b': torch.tensor([4, 5, 6])},
                {'a': torch.tensor([7, 8, 9]), 'b': torch.tensor([10, 11, 12])}
            ]
        >>> collated = ttorch_collate(nested_dicts)
        collated = {
            'a': torch.tensor([[1, 2, 3], [7, 8, 9]]),
            'b': torch.tensor([[4, 5, 6], [10, 11, 12]])
        }
    """

    def inplace_fn(t):
        for k in t.keys():
            if isinstance(t[k], torch.Tensor):
                if len(t[k].shape) == 2 and t[k].shape[1] == 1:  # reshape (B, 1) -> (B)
                    t[k] = t[k].squeeze(-1)
            else:
                inplace_fn(t[k])

    x = ttorch.stack(x)
    if cat_1dim:
        inplace_fn(x)
    if json:
        x = x.json()
    return x


def default_collate(batch: Sequence,
                    cat_1dim: bool = True,
                    ignore_prefix: list = ['collate_ignore']) -> Union[torch.Tensor, Mapping, Sequence]:
    """
    Overview:
        Put each data field into a tensor with outer dimension batch size.

    Arguments:
        - batch (:obj:`Sequence`): A data sequence, whose length is batch size, whose element is one piece of data.
        - cat_1dim (:obj:`bool`): Whether to concatenate tensors with shape (B, 1) to (B), defaults to True.
        - ignore_prefix (:obj:`list`): A list of prefixes to ignore when collating dictionaries, \
            defaults to ['collate_ignore'].

    Returns:
        - ret (:obj:`Union[torch.Tensor, Mapping, Sequence]`): the collated data, with batch size into each data \
            field. The return dtype depends on the original element dtype, can be [torch.Tensor, Mapping, Sequence].

    Example:
        >>> # a list with B tensors shaped (m, n) -->> a tensor shaped (B, m, n)
        >>> a = [torch.zeros(2,3) for _ in range(4)]
        >>> default_collate(a).shape
        torch.Size([4, 2, 3])
        >>>
        >>> # a list with B lists, each list contains m elements -->> a list of m tensors, each with shape (B, )
        >>> a = [[0 for __ in range(3)] for _ in range(4)]
        >>> default_collate(a)
        [tensor([0, 0, 0, 0]), tensor([0, 0, 0, 0]), tensor([0, 0, 0, 0])]
        >>>
        >>> # a list with B dicts, whose values are tensors shaped :math:`(m, n)` -->>
        >>> # a dict whose values are tensors with shape :math:`(B, m, n)`
        >>> a = [{i: torch.zeros(i,i+1) for i in range(2, 4)} for _ in range(4)]
        >>> print(a[0][2].shape, a[0][3].shape)
        torch.Size([2, 3]) torch.Size([3, 4])
        >>> b = default_collate(a)
        >>> print(b[2].shape, b[3].shape)
        torch.Size([4, 2, 3]) torch.Size([4, 3, 4])
    """

    if isinstance(batch, ttorch.Tensor):
        return batch.json()

    elem = batch[0]
    elem_type = type(elem)
    if isinstance(elem, torch.Tensor):
        out = None
        if torch_ge_131() and torch.utils.data.get_worker_info() is not None:
            # If we're in a background process, directly concatenate into a
            # shared memory tensor to avoid an extra copy
            numel = sum([x.numel() for x in batch])
            storage = elem.storage()._new_shared(numel)
            out = elem.new(storage)
        if elem.shape == (1, ) and cat_1dim:
            # reshape (B, 1) -> (B)
            return torch.cat(batch, 0, out=out)
            # return torch.stack(batch, 0, out=out)
        else:
            return torch.stack(batch, 0, out=out)
    elif isinstance(elem, ttorch.Tensor):
        return ttorch_collate(batch, json=True, cat_1dim=cat_1dim)
    elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \
            and elem_type.__name__ != 'string_':
        if elem_type.__name__ == 'ndarray':
            # array of string classes and object
            if np_str_obj_array_pattern.search(elem.dtype.str) is not None:
                raise TypeError(default_collate_err_msg_format.format(elem.dtype))
            return default_collate([torch.as_tensor(b) for b in batch], cat_1dim=cat_1dim)
        elif elem.shape == ():  # scalars
            return torch.as_tensor(batch)
    elif isinstance(elem, float):
        return torch.tensor(batch, dtype=torch.float32)
    elif isinstance(elem, int_classes):
        dtype = torch.bool if isinstance(elem, bool) else torch.int64
        return torch.tensor(batch, dtype=dtype)
    elif isinstance(elem, string_classes):
        return batch
    elif isinstance(elem, container_abcs.Mapping):
        ret = {}
        for key in elem:
            if any([key.startswith(t) for t in ignore_prefix]):
                ret[key] = [d[key] for d in batch]
            else:
                ret[key] = default_collate([d[key] for d in batch], cat_1dim=cat_1dim)
        return ret
    elif isinstance(elem, tuple) and hasattr(elem, '_fields'):  # namedtuple
        return elem_type(*(default_collate(samples, cat_1dim=cat_1dim) for samples in zip(*batch)))
    elif isinstance(elem, container_abcs.Sequence):
        transposed = zip(*batch)
        return [default_collate(samples, cat_1dim=cat_1dim) for samples in transposed]

    raise TypeError(default_collate_err_msg_format.format(elem_type))


def timestep_collate(batch: List[Dict[str, Any]]) -> Dict[str, Union[torch.Tensor, list]]:
    """
    Overview:
        Collates a batch of timestepped data fields into tensors with the outer dimension being the batch size. \
        Each timestepped data field is represented as a tensor with shape [T, B, any_dims], where T is the length \
        of the sequence, B is the batch size, and any_dims represents the shape of the tensor at each timestep.

    Arguments:
        - batch(:obj:`List[Dict[str, Any]]`): A list of dictionaries with length B, where each dictionary represents \
            a timestepped data field. Each dictionary contains a key-value pair, where the key is the name of the \
            data field and the value is a sequence of torch.Tensor objects with any shape.

    Returns:
        - ret(:obj:`Dict[str, Union[torch.Tensor, list]]`): The collated data, with the timestep and batch size \
            incorporated into each data field. The shape of each data field is [T, B, dim1, dim2, ...].

    Examples:
        >>> batch = [
                {'data0': [torch.tensor([1, 2, 3]), torch.tensor([4, 5, 6])]},
                {'data1': [torch.tensor([7, 8, 9]), torch.tensor([10, 11, 12])]}
            ]
        >>> collated_data = timestep_collate(batch)
        >>> print(collated_data['data'].shape)
        torch.Size([2, 2, 3])
    """

    def stack(data):
        if isinstance(data, container_abcs.Mapping):
            return {k: stack(data[k]) for k in data}
        elif isinstance(data, container_abcs.Sequence) and isinstance(data[0], torch.Tensor):
            return torch.stack(data)
        else:
            return data

    elem = batch[0]
    assert isinstance(elem, (container_abcs.Mapping, list)), type(elem)
    if isinstance(batch[0], list):  # new pipeline + treetensor
        prev_state = [[b[i].get('prev_state') for b in batch] for i in range(len(batch[0]))]
        batch_data = ttorch.stack([ttorch_collate(b) for b in batch])  # (B, T, *)
        del batch_data.prev_state
        batch_data = batch_data.transpose(1, 0)
        batch_data.prev_state = prev_state
    else:
        prev_state = [b.pop('prev_state') for b in batch]
        batch_data = default_collate(batch)  # -> {some_key: T lists}, each list is [B, some_dim]
        batch_data = stack(batch_data)  # -> {some_key: [T, B, some_dim]}
        transformed_prev_state = list(zip(*prev_state))
        batch_data['prev_state'] = transformed_prev_state
        # append back prev_state, avoiding multi batch share the same data bug
        for i in range(len(batch)):
            batch[i]['prev_state'] = prev_state[i]
    return batch_data


def diff_shape_collate(batch: Sequence) -> Union[torch.Tensor, Mapping, Sequence]:
    """
    Overview:
        Collates a batch of data with different shapes.
        This function is similar to `default_collate`, but it allows tensors in the batch to have `None` values, \
        which is common in StarCraft observations.

    Arguments:
        - batch (:obj:`Sequence`): A sequence of data, where each element is a piece of data.

    Returns:
        - ret (:obj:`Union[torch.Tensor, Mapping, Sequence]`): The collated data, with the batch size applied \
            to each data field. The return type depends on the original element type and can be a torch.Tensor, \
            Mapping, or Sequence.

    Examples:
        >>> # a list with B tensors shaped (m, n) -->> a tensor shaped (B, m, n)
        >>> a = [torch.zeros(2,3) for _ in range(4)]
        >>> diff_shape_collate(a).shape
        torch.Size([4, 2, 3])
        >>>
        >>> # a list with B lists, each list contains m elements -->> a list of m tensors, each with shape (B, )
        >>> a = [[0 for __ in range(3)] for _ in range(4)]
        >>> diff_shape_collate(a)
        [tensor([0, 0, 0, 0]), tensor([0, 0, 0, 0]), tensor([0, 0, 0, 0])]
        >>>
        >>> # a list with B dicts, whose values are tensors shaped :math:`(m, n)` -->>
        >>> # a dict whose values are tensors with shape :math:`(B, m, n)`
        >>> a = [{i: torch.zeros(i,i+1) for i in range(2, 4)} for _ in range(4)]
        >>> print(a[0][2].shape, a[0][3].shape)
        torch.Size([2, 3]) torch.Size([3, 4])
        >>> b = diff_shape_collate(a)
        >>> print(b[2].shape, b[3].shape)
        torch.Size([4, 2, 3]) torch.Size([4, 3, 4])
    """
    elem = batch[0]
    elem_type = type(elem)
    if any([isinstance(elem, type(None)) for elem in batch]):
        return batch
    elif isinstance(elem, torch.Tensor):
        shapes = [e.shape for e in batch]
        if len(set(shapes)) != 1:
            return batch
        else:
            return torch.stack(batch, 0)
    elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \
            and elem_type.__name__ != 'string_':
        if elem_type.__name__ == 'ndarray':
            return diff_shape_collate([torch.as_tensor(b) for b in batch])  # todo
        elif elem.shape == ():  # scalars
            return torch.as_tensor(batch)
    elif isinstance(elem, float):
        return torch.tensor(batch, dtype=torch.float32)
    elif isinstance(elem, int_classes):
        dtype = torch.bool if isinstance(elem, bool) else torch.int64
        return torch.tensor(batch, dtype=dtype)
    elif isinstance(elem, Mapping):
        return {key: diff_shape_collate([d[key] for d in batch]) for key in elem}
    elif isinstance(elem, tuple) and hasattr(elem, '_fields'):  # namedtuple
        return elem_type(*(diff_shape_collate(samples) for samples in zip(*batch)))
    elif isinstance(elem, Sequence):
        transposed = zip(*batch)
        return [diff_shape_collate(samples) for samples in transposed]

    raise TypeError('not support element type: {}'.format(elem_type))


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]`): The collated data batch. It can be a tensor, \
            sequence, or mapping.
        - ignore(:obj:`List[str]`): A list of names to be ignored. Only applicable if the input ``batch`` is a \
            dictionary. If a key is in this list, its value will remain the same without decollation. Defaults to \
            ['prev_state', 'prev_actor_state', 'prev_critic_state'].

    Returns:
        - ret (:obj:`List[Any]`): A list with B elements, where B is the batch size.

    Examples:
        >>> batch = {
            'a': [
                [1, 2, 3],
                [4, 5, 6]
            ],
            'b': [
                [7, 8, 9],
                [10, 11, 12]
            ]}
        >>> default_decollate(batch)
        {
            0: {'a': [1, 2, 3], 'b': [7, 8, 9]},
            1: {'a': [4, 5, 6], 'b': [10, 11, 12]},
        }
    """
    if isinstance(batch, torch.Tensor):
        batch = torch.split(batch, 1, dim=0)
        # Squeeze if the 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)]
    elif isinstance(batch, torch.distributions.Distribution):  # For compatibility
        return [None for _ in range(batch.batch_shape[0])]

    raise TypeError("Not supported batch type: {}".format(type(batch)))