import torch


def prepare_loss_weights(labels, pos_cls_weight=1.0, neg_cls_weight=1.0, dtype=torch.float32):
    """
    get cls_weights and reg_weights from labels.
    :param labels: ground truth labels
    :param pos_cls_weight: positive classification weight
    :param neg_cls_weight: negative classification weight
    :param dtype: the type of data
    :return: the weights of classification and localization with the selected labels
    """
    cared = labels >= 0
    # cared: [N, num_anchors]
    positives = labels > 0
    negatives = labels == 0
    negative_cls_weights = negatives.type(dtype) * neg_cls_weight
    cls_weights = negative_cls_weights + pos_cls_weight * positives.type(dtype)
    reg_weights = positives.type(dtype)

    pos_normalizer = positives.sum(1, keepdim=True).type(dtype)
    reg_weights /= torch.clamp(pos_normalizer, min=1.0)
    cls_weights /= torch.clamp(pos_normalizer, min=1.0)

    return cls_weights, reg_weights, cared