import random
from typing import Callable, Union, List

from ding.data.buffer import BufferedData
from ding.utils import fastcopy


def padding(policy="random"):
    """
    Overview:
        Fill the nested buffer list to the same size as the largest list.
        The default policy `random` will randomly select data from each group
        and fill it into the current group list.
    Arguments:
        - policy (:obj:`str`): Padding policy, supports `random`, `none`.
    """

    def sample(chain: Callable, *args, **kwargs) -> Union[List[BufferedData], List[List[BufferedData]]]:
        sampled_data = chain(*args, **kwargs)
        if len(sampled_data) == 0 or isinstance(sampled_data[0], BufferedData):
            return sampled_data
        max_len = len(max(sampled_data, key=len))
        for i, grouped_data in enumerate(sampled_data):
            group_len = len(grouped_data)
            if group_len == max_len:
                continue
            for _ in range(max_len - group_len):
                if policy == "random":
                    sampled_data[i].append(fastcopy.copy(random.choice(grouped_data)))
                elif policy == "none":
                    sampled_data[i].append(BufferedData(data=None, index=None, meta=None))

        return sampled_data

    def _padding(action: str, chain: Callable, *args, **kwargs):
        if action == "sample":
            return sample(chain, *args, **kwargs)
        return chain(*args, **kwargs)

    return _padding