# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/transforms/04_clipping.ipynb.

# %% auto 0
__all__ = ['Clip', 'Clipd']

# %% ../../nbs/transforms/04_clipping.ipynb 2
from collections.abc import Hashable

import torch
from monai.config import KeysCollection
from monai.transforms import Transform

# %% ../../nbs/transforms/04_clipping.ipynb 4
class Clip(Transform):
    def __init__(self, min_value: float, max_value: float):
        super().__init__()

        self.min_value = min_value
        self.max_value = max_value

    def __call__(self, data: torch.Tensor) -> torch.Tensor:
        """
        Clip the input tensor to the specified range.

        Args:
            data: Input tensor to be clipped.

        Returns:
            Clipped tensor.
        """
        return torch.clamp(data, self.min_value, self.max_value)

    def __repr__(self) -> str:
        return f"Clip(min_value={self.min_value}, max_value={self.max_value})"

# %% ../../nbs/transforms/04_clipping.ipynb 6
class Clipd(Transform):
    def __init__(self, keys: KeysCollection, min_value: float, max_value: float):
        super().__init__()

        self.keys = keys
        self.transform = Clip(min_value, max_value)

    def __call__(self, data: dict[Hashable, torch.Tensor]) -> dict[Hashable, torch.Tensor]:
        """
        Clip the input tensor to the specified range.

        Args:
            data: Input tensor to be clipped.

        Returns:
            Clipped tensor.
        """
        for key in self.keys:
            if key in data:
                data[key] = self.transform(data[key])
            else:
                raise KeyError(f"Key {key} not found in input data.")

        return data

    def __repr__(self) -> str:
        return f"Clipd(keys={self.keys}, min_value={self.transform.min_value}, max_value={self.transform.max_value})"
