import numpy as np
import torch

from funasr.datasets.large_datasets.collate_fn import crop_to_max_size


def clipping(data):
    assert isinstance(data, list)
    assert "key" in data[0]

    keys = [x["key"] for x in data]

    batch = {}
    data_names = data[0].keys()
    for data_name in data_names:
        if data_name == "key":
            continue
        else:
            if data[0][data_name].dtype.kind == "i":
                tensor_type = torch.int64
            else:
                tensor_type = torch.float32

            tensor_list = [torch.tensor(np.copy(d[data_name]), dtype=tensor_type) for d in data]
            tensor_lengths = torch.tensor([len(d[data_name]) for d in data], dtype=torch.int32)

            length_clip = min(tensor_lengths)
            tensor_clip = tensor_list[0].new_zeros(
                len(tensor_list), length_clip, tensor_list[0].shape[1]
            )
            for i, (tensor, length) in enumerate(zip(tensor_list, tensor_lengths)):
                diff = length - length_clip
                assert diff >= 0
                if diff == 0:
                    tensor_clip[i] = tensor
                else:
                    tensor_clip[i] = crop_to_max_size(tensor, length_clip)

            batch[data_name] = tensor_clip
            batch[data_name + "_lengths"] = torch.tensor(
                [tensor.shape[0] for tensor in tensor_clip], dtype=torch.long
            )

    return keys, batch