import numpy as np


class RandomCrop:
    """
    Crop the 2D numpy [c, h, w] and mask [h, w] in a sample
    """

    def __init__(self, output_size, training=True):
        assert isinstance(output_size, (int, tuple))
        if isinstance(output_size, int):
            self.output_size = (output_size, output_size)
        else:
            assert len(output_size) == 2
            self.output_size = output_size

        self.training = training

    def __call__(self, sample):
        image, mask = sample['image'], sample['mask']

        c, h, w = image.shape
        new_h, new_w = self.output_size

        assert (h - new_h) >= 0 and (w - new_w) >= 0, f"The orig size:({h},{w}) should be equal or larger than " \
                                                      f"crop size:({new_h},{new_w})"

        if not self.training:
            # Test stage, will set the crop position at 0,0
            top = 0
            left = 0
        else:
            top = np.random.randint(h - new_h + 1)
            left = np.random.randint(w - new_w + 1)

        image = image[:, top: top + new_h, left: left + new_w]
        mask = mask[top: top + new_h, left: left + new_w]

        return {'image': image, 'mask': mask}
