import torch


def psnr(gt, pred, data_range=None, batch=True, reduce=True):
    """ Compute the peak signal to noise ratio (psnr)
    :param gt: gt image (torch.Tensor
    :param pred: input image (torch.Tensor)
    :param data_range: if None, estimated from gt
    :return: (mean) psnr
    """
    if batch:
        batch_size = gt.shape[0]
    else:
        batch_size = 1

    # reshape the view
    pred = pred.contiguous().view(batch_size, -1)
    gt = gt.contiguous().view(batch_size, -1)

    if data_range is None:
        # by default use max, same as fastmri
        data_range = gt.max(dim=1)[0]# - gt.min(dim=1)[0]

    mse_err = (abs(gt - pred) ** 2).mean(1)
    psnr_val = 10 * torch.log10(data_range ** 2 / mse_err)
    if reduce:
        return psnr_val.mean()
    else:
        return psnr_val