OneOf

class OneOf(Transform):
    """Apply one of the given transforms, chosen at random.

    When applied to a batch with `per_instance=True` (the default),
    each batch element independently chooses which transform to apply.
    This requires shape- and schema-preserving transforms so the
    elements can be re-stacked. Pass `per_instance=False` to choose a
    single transform for the whole batch.

    Args:
        transforms: Sequence of transforms, or a `dict` mapping
            transforms to their relative weights. If a sequence is
            given, all transforms have equal probability.
        **kwargs: See [`Transform`][torchio.Transform] for additional
            keyword arguments.

    Examples:
        >>> import torchio as tio
        >>> augmentation = tio.OneOf({
        ...     tio.Noise(std=0.1): 0.7,
        ...     tio.Flip(axes=(0,)): 0.3,
        ... })
    """

    def __init__(
        self,
        transforms: Sequence[Transform] | dict[Transform, float],
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        if isinstance(transforms, dict):
            weight_dict = cast(dict[Transform, float], transforms)
            self.transforms = list(weight_dict.keys())
            w_list: list[float] = list(weight_dict.values())
            total: float = sum(w_list)
            self.weights = [w / total for w in w_list]
        else:
            self.transforms = list(transforms)
            n = len(self.transforms)
            self.weights = [1.0 / n] * n

    def forward(self, data):
        if self.copy:
            data = copy.deepcopy(data)
        batch, unwrap = self._wrap(data)
        # The input is copied once above, so children apply without copying.
        with _disabled_copy(self.transforms):
            if self.per_instance and batch.batch_size > 1:
                return unwrap(self._forward_per_element(batch))
            if torch.rand(1).item() >= self.p:
                return unwrap(batch)
            idx = int(
                torch.multinomial(
                    torch.tensor(self.weights),
                    num_samples=1,
                ).item()
            )
            batch = self.transforms[idx](batch)
            return unwrap(batch)

    def _forward_per_element(self, batch):
        """Apply an independently chosen transform to each batch element."""
        if self.p == 0:
            return batch
        weights = torch.tensor(self.weights)
        out_subjects = []
        any_applied = False
        for subject in batch.unbatch():
            if torch.rand(1).item() < self.p:
                any_applied = True
                idx = int(torch.multinomial(weights, num_samples=1).item())
                subject = _apply_to_element(subject, self.transforms[idx])
            out_subjects.append(subject)
        if not any_applied:
            return batch
        return _rebatch_with_history(out_subjects, "OneOf")

    def to_hydra(self) -> dict[str, Any]:
        cfg = super().to_hydra()
        cfg["transforms"] = [t.to_hydra() for t in self.transforms]
        return cfg