# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/transforms/02_croppad.ipynb.

# %% auto 0
__all__ = ['get_updated_crop_start', 'CropForegroundWithCropTrackingd', 'RandSpatialCropSamplesWithCropTracking',
           'RandSpatialCropSamplesWithCropTrackingd']

# %% ../../nbs/transforms/02_croppad.ipynb 2
from collections.abc import Sequence

import torch
from monai.data import MetaTensor
from monai.data.meta_obj import get_track_meta
from monai.transforms.croppad.array import RandSpatialCropSamples
from monai.transforms.croppad.dictionary import CropForegroundd, RandSpatialCropSamplesd
from monai.utils import ImageMetaKey as Key

# %% ../../nbs/transforms/02_croppad.ipynb 5
def get_updated_crop_start(current_crop_start, new_crop_start):
    if not torch.is_tensor(new_crop_start):
        new_crop_start = torch.tensor(new_crop_start)

    if current_crop_start is None:
        return new_crop_start

    if not torch.is_tensor(current_crop_start):
        current_crop_start = torch.tensor(current_crop_start)

    updated_crop_start = current_crop_start + new_crop_start
    return updated_crop_start

# %% ../../nbs/transforms/02_croppad.ipynb 9
class CropForegroundWithCropTrackingd(CropForegroundd):
    def __init__(
        self,
        crop_offset_key: str = "crop_offset",
        *args,
        **kwargs,
    ) -> MetaTensor:
        super().__init__(*args, **kwargs)
        self.crop_offset_key = crop_offset_key

    def __call__(self, data, *args, **kwargs):
        output = super().__call__(data, *args, **kwargs)
        crop_offset = output[self.start_coord_key]
        output[self.crop_offset_key] = get_updated_crop_start(output.get(self.crop_offset_key), crop_offset)
        return output

# %% ../../nbs/transforms/02_croppad.ipynb 11
class RandSpatialCropSamplesWithCropTracking(RandSpatialCropSamples):  # To return the crops along with the crop offset
    def __init__(self, crop_offset_key: str = "crop_offset", *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.crop_key = crop_offset_key

    def __call__(self, img: torch.Tensor, lazy: bool | None = None) -> list[torch.Tensor]:
        """
        Apply the transform to `img`, assuming `img` is channel-first and
        cropping doesn't change the channel dim.
        """
        ret = []
        lazy_ = self.lazy if lazy is None else lazy
        for i in range(self.num_samples):
            cropped = self.cropper(img, lazy=lazy_)
            if get_track_meta():
                cropped.meta[self.crop_key] = tuple(_slice.start for _slice in self.cropper._slices)
                cropped.meta[Key.PATCH_INDEX] = i  # type: ignore
                self.push_transform(cropped, replace=True, lazy=lazy_)  # track as this class instead of RandSpatialCrop
            ret.append(cropped)
        return ret


class RandSpatialCropSamplesWithCropTrackingd(RandSpatialCropSamplesd):
    def __init__(
        self,
        keys,
        roi_size: Sequence[int] | int,
        num_samples: int,
        max_roi_size: Sequence[int] | int | None = None,
        random_center: bool = True,
        random_size: bool = False,
        allow_missing_keys: bool = False,
        lazy: bool = False,
        crop_offset_key: str = "crop_offset",
    ) -> MetaTensor:
        super().__init__(
            keys, roi_size, num_samples, max_roi_size, random_center, random_size, allow_missing_keys, lazy
        )
        self.crop_offset_key = crop_offset_key
        self.cropper = RandSpatialCropSamplesWithCropTracking(
            crop_offset_key, roi_size, num_samples, max_roi_size, random_center, random_size, lazy=lazy
        )

    def __call__(self, data, *args, **kwargs):
        output = super().__call__(data, *args, **kwargs)
        for key in self.keys:
            for o in output:
                crop_offset = o[key].meta.get(self.crop_offset_key)
                o[self.crop_offset_key] = get_updated_crop_start(o.get(self.crop_offset_key), crop_offset)
        return output
