# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/transforms/03_spatial.ipynb.

# %% auto 0
__all__ = ['RandFlipWithCropTrackingd', 'ResizedWithCropTrackingd']

# %% ../../nbs/transforms/03_spatial.ipynb 2
import torch
from monai.data import MetaTensor
from monai.transforms.spatial.dictionary import RandFlipd, Resized

# %% ../../nbs/transforms/03_spatial.ipynb 4
class RandFlipWithCropTrackingd(RandFlipd):
    def __init__(
        self,
        crop_offset_key: str = "crop_offset",
        original_shape_key: str = "original_shape",
        *args,
        **kwargs,
    ) -> MetaTensor:
        super().__init__(*args, **kwargs)
        self.crop_offset_key = crop_offset_key
        self.original_shape_key = original_shape_key

    def __call__(self, data, *args, **kwargs):
        output = super().__call__(data, *args, **kwargs)
        # Check if image was flipped and was cropped previously
        if self._do_transform and self.crop_offset_key in data:
            # Ensure all inputs had input of same size for consistent crop value
            shapes = [data[key].shape for key in self.keys]
            if not all(shape == shapes[0] for shape in shapes):
                raise ValueError("All inputs must have the same shape for consistent crop value")
            crop_shape = shapes[0]
            orig_shape = torch.Tensor(data[self.original_shape_key])
            # New crop value will be flipped
            crop_offset = torch.Tensor(data[self.crop_offset_key])
            # Flip crop value
            crop_offset[-1] = orig_shape[-1] - crop_offset[-1] - crop_shape[-1]
            # Update crop value
            output[self.crop_offset_key] = crop_offset
        return output

# %% ../../nbs/transforms/03_spatial.ipynb 6
class ResizedWithCropTrackingd(Resized):
    def __init__(
        self,
        crop_offset_key: str = "crop_offset",
        original_shape_key: str = "original_shape",
        *args,
        **kwargs,
    ) -> MetaTensor:
        super().__init__(*args, **kwargs)
        self.crop_offset_key = crop_offset_key
        self.original_shape_key = original_shape_key

    def __call__(self, data, *args, **kwargs):
        # Check if image was cropped previously
        if self.crop_offset_key in data:
            # Ensure all inputs had input of same size for consistent crop value
            shapes = [data[key].shape for key in self.keys]
            if not all(shape == shapes[0] for shape in shapes):
                raise ValueError("All inputs must have the same shape for consistent crop value")
            old_shape = data[self.keys[0]].shape[1:]

        output = super().__call__(data, *args, **kwargs)

        # Check if image was cropped previously
        if self.crop_offset_key in data:
            # Calculate scale of resizing along every axis
            new_shape = output[self.keys[0]].shape[1:]
            scale = torch.tensor(new_shape) / torch.tensor(old_shape)

            # Update crop_offset
            crop_offset = torch.Tensor(data[self.crop_offset_key])
            output[self.crop_offset_key] = (crop_offset * scale).int()

            # Update original_shape
            if self.original_shape_key in data:
                original_shape = torch.Tensor(data[self.original_shape_key])
                output[self.original_shape_key] = (original_shape * scale).int()

        return output
