import torch


def generate_softmax_labels(rois, keyps, label_h, label_w):
    """
    Arguments:
        rois (FloatTensor): [R, 4]
        keyps (FloatTensor): [R, K, 3], (x, y, v)
        label_w, label_h (int): label size

    Returns:
        labels (LongTensor): [R, K]
    """
    assert (rois.shape[0] == keyps.shape[0])
    R, K = keyps.shape[:2]
    scale_x = label_w / torch.clamp(rois[:, 2] - rois[:, 0] + 1, min=1)
    scale_y = label_h / torch.clamp(rois[:, 3] - rois[:, 1] + 1, min=1)
    keyp_x = (keyps[:, :, 0] - rois[:, 0, None] + 0.5) * scale_x[:, None]
    keyp_y = (keyps[:, :, 1] - rois[:, 1, None] + 0.5) * scale_y[:, None]
    keyp_x = keyp_x.to(dtype=torch.int64)
    keyp_y = keyp_y.to(dtype=torch.int64)
    keyp_v = keyps[:, :, 2] > 0

    mask = ((keyp_x >= 0) & (keyp_x < label_w) & (keyp_y >= 0) & (keyp_y < label_h) & keyp_v)
    labels = -torch.ones((R, K), device=rois.device, dtype=torch.int64)
    labels[mask] = keyp_y[mask] * label_w + keyp_x[mask]
    return labels