import torch


def lengths_to_mask(lengths):
    max_len = max(lengths)
    mask = torch.arange(max_len, device=lengths.device).expand(
        len(lengths), max_len) < lengths.unsqueeze(1)
    return mask


def collate_tensors(batch):
    dims = batch[0].dim()
    max_size = [max([b.size(i) for b in batch]) for i in range(dims)]
    size = (len(batch),) + tuple(max_size)
    canvas = batch[0].new_zeros(size=size)
    for i, b in enumerate(batch):
        sub_tensor = canvas[i]
        for d in range(dims):
            sub_tensor = sub_tensor.narrow(d, 0, b.size(d))
        sub_tensor.add_(b)
    return canvas


def collate(batch):
    databatch = [b[0] for b in batch]
    labelbatch = [b[1] for b in batch]
    lenbatch = [len(b[0][0][0]) for b in batch]

    databatchTensor = collate_tensors(databatch)
    labelbatchTensor = torch.as_tensor(labelbatch)
    lenbatchTensor = torch.as_tensor(lenbatch)

    maskbatchTensor = lengths_to_mask(lenbatchTensor)
    # x - [bs, njoints, nfeats, lengths]
    #   - nfeats, the representation of a joint
    # y - [bs]
    # mask - [bs, lengths]
    # lengths - [bs]
    batch = {"x": databatchTensor, "y": labelbatchTensor,
             "mask": maskbatchTensor, 'lengths': lenbatchTensor}
    return batch


# slow version with padding
def collate_data3d_slow(batch):
    batchTensor = {}
    for key in batch[0].keys():
        databatch = [b[key] for b in batch]
        batchTensor[key] = collate_tensors(databatch)
    batch = batchTensor
    # theta - [bs, lengths, 85], theta shape (85,)
    #       - (np.array([1., 0., 0.]), pose(72), shape(10)), axis=0)
    # kp_2d - [bs, lengths, njoints, nfeats], nfeats (x,y,weight)
    # kp_3d - [bs, lengths, njoints, nfeats], nfeats (x,y,z)
    # w_smpl - [bs, lengths] zeros
    # w_3d - [bs, lengths] zeros
    return batch

def collate_data3d(batch):
    batchTensor = {}
    for key in batch[0].keys():
        databatch = [b[key] for b in batch]
        if key == "paths":
            batchTensor[key] = databatch
        else:    
            batchTensor[key] = torch.stack(databatch,axis=0)
    batch = batchTensor
    # theta - [bs, lengths, 85], theta shape (85,)
    #       - (np.array([1., 0., 0.]), pose(72), shape(10)), axis=0)
    # kp_2d - [bs, lengths, njoints, nfeats], nfeats (x,y,weight)
    # kp_3d - [bs, lengths, njoints, nfeats], nfeats (x,y,z)
    # w_smpl - [bs, lengths] zeros
    # w_3d - [bs, lengths] zeros
    return batch