import torch


def compute_boxes(scores, offsets, scale, threshold):
    """Return bounding boxes, scores, offsets, and batch indices in matrix.

    PNet acts like a 12x12 convolution with stride 2,
    so need to convert bounding box indices back to original image coordinates.

    Arguments
    ---------
    scores : torch.Tensor
        size [n, 1, h, w]
        score for face presence at each image location

    offsets : torch.Tensor
        size[n, 4, h, w]
        offsets for each image location to recover full image coordinates

    scale : float
        scaling of original image prior to PNet application

    threshold : float
        minimum score value for inclusion

    Returns
    -------
    bounding_boxes : torch.Tensor
        size [num_boxes, 10]
        Each row is a single bounding box.

        Column 0 is batch index.
        Columns 1 - 4 are bounding box top left and bottom right coordinates.
        Column 5 is score for that box.
        Columns 6-10 are offset values.
    """
    stride = 2
    kernel_size = 12

    detection_indices = (scores > threshold).nonzero(as_tuple=True)
    batch_ix = detection_indices[0]

    if batch_ix.size()[0] == 0:
        return None

    offsets_ = offsets[batch_ix, :, detection_indices[1], detection_indices[2]]
    h_ix, w_ix = [stride * d + 1 for d in detection_indices[1:]]
    scores_ = scores[batch_ix, detection_indices[1], detection_indices[2]]

    bounding_boxes = torch.stack([batch_ix.to(torch.float32),
                                  torch.round(w_ix / scale),
                                  torch.round(h_ix / scale),
                                  torch.round((w_ix + kernel_size) / scale),
                                  torch.round((h_ix + kernel_size) / scale),
                                  scores_,
                                  offsets_[:, 0],
                                  offsets_[:, 1],
                                  offsets_[:, 2],
                                  offsets_[:, 3]], dim=1)

    return bounding_boxes