import torch


def get_center_block_mask(samples: torch.Tensor, mask_size: int, overlap: int):
    """
    Mask out the center square region of samples provided as input.

    Parameters
    ----------
    samples: torch.Tensor
        Batch of samples, e.g. images, which are passed through the network and for which specified intermediate
        results are extracted
    mask_size: int
        Size of the squared block mask in pixels.
    overlap: int
        Specify number of pixels of overlap. The reconstruction loss puts a larger weight on "overlap"-pixels to
        force the generator to inpaint the masked part of the image with high consistenncy at the borders.
    Returns
    -------
    masked_samples: torch.Tensor
        Tensor containing samples to which the random mask is applied.
    true_masked_part: torch.Tensor
        Tensor containing the original image pixels, which have been masked out.
    mask_indices: Tuple
        Tuple which contains x and y indices of the upper left corner of the region, which has been masked out.
    """
    img_size = samples.size()[-1]
    center_index = (img_size - mask_size) // 2
    masked_samples = samples.clone()

    # Image is not masked out in overlap region
    m1, m2 = center_index + overlap, center_index + mask_size - overlap

    masked_samples[:, 0, m1:m2, m1:m2] = 2 * 117.0 / 255.0 - 1.0
    masked_samples[:, 1, m1:m2, m1:m2] = 2 * 104.0 / 255.0 - 1.0
    masked_samples[:, 2, m1:m2, m1:m2] = 2 * 123.0 / 255.0 - 1.0

    true_masked_part = samples[:, :, center_index:center_index+mask_size, center_index:center_index+mask_size]
    return masked_samples, true_masked_part, (center_index, center_index)