"""Module containing base interfaces for all transform implementations. of alb

This module defines the fundamental transform interfaces that form the base hierarchy for
all transformation classes in Albumentations. It provides abstract classes and mixins that
define common behavior for image, keypoint, bounding box, and volumetric transformations.
The interfaces handle parameter validation, random state management, target type checking,
and serialization capabilities that are inherited by concrete transform implementations.
"""

import inspect
import random
from collections.abc import Callable, Sequence
from copy import deepcopy
from typing import Any
from warnings import warn

import cv2
import numpy as np
from albucore import batch_transform
from pydantic import BaseModel, ConfigDict, Field

from albumentations.core.bbox_utils import BboxProcessor
from albumentations.core.keypoints_utils import KeypointsProcessor
from albumentations.core.validation import ValidatedTransformMeta

from .serialization import Serializable, SerializableMeta, get_shortest_class_fullname
from .type_definitions import ALL_TARGETS, ImageType, Targets, VolumeType
from .utils import format_args

__all__ = ["BasicTransform", "CustomTransformsApplyMixin", "DualTransform", "ImageOnlyTransform", "NoOp", "Transform3D"]


class Interpolation:
    def __init__(self, downscale: int = cv2.INTER_NEAREST, upscale: int = cv2.INTER_NEAREST):
        self.downscale = downscale
        self.upscale = upscale


class BaseTransformInitSchema(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)
    p: float = Field(ge=0, le=1)
    strict: bool


class CombinedMeta(SerializableMeta, ValidatedTransformMeta):
    pass


class BasicTransform(Serializable, metaclass=CombinedMeta):
    """Base class for all transforms in Albumentations. Provides core functionality for application,
    serialization, and params.

    This class provides core functionality for transform application, serialization,
    and parameter handling. It defines the interface that all transforms must follow
    and implements common methods used across different transform types.

    Class Attributes:
        _targets (tuple[Targets, ...] | Targets): Target types this transform can work with.
        _available_keys (set[str]): String representations of valid target keys.
        _key2func (dict[str, Callable[..., Any]]): Mapping between target keys and their processing functions.

    Args:
        interpolation (int): Interpolation method for image transforms.
        fill (int | float | list[int] | list[float]): Fill value for image padding.
        fill_mask (int | float | list[int] | list[float]): Fill value for mask padding.
        deterministic (bool, optional): Whether the transform is deterministic.
        save_key (str, optional): Key for saving transform parameters.
        replay_mode (bool, optional): Whether the transform is in replay mode.
        applied_in_replay (bool, optional): Whether the transform was applied in replay.
        p (float): Probability of applying the transform.

    Note:
        The base class methods use *args to allow subclasses to add specific named parameters
        (e.g., def apply(self, img, gamma, **params) is a valid override of apply(self, img, *args, **params)).

    """

    _targets: tuple[Targets, ...] | Targets  # targets that this transform can work on
    _available_keys: set[str]  # targets that this transform, as string, lower-cased
    _key2func: dict[
        str,
        Callable[..., Any],
    ]  # mapping for targets (plus additional targets) and methods for which they depend
    call_backup = None
    interpolation: int
    fill: tuple[float, ...] | float
    fill_mask: tuple[float, ...] | float | None
    # replay mode params
    deterministic: bool = False
    save_key = "replay"
    replay_mode = False
    applied_in_replay = False

    class InitSchema(BaseTransformInitSchema):
        pass

    def __init__(self, p: float = 0.5):
        self.p = p
        self._additional_targets: dict[str, str] = {}
        self.params: dict[Any, Any] = {}
        self._key2func = {}
        self._set_keys()
        self.processors: dict[str, BboxProcessor | KeypointsProcessor] = {}
        self.seed: int | None = None
        self.set_random_seed(self.seed)
        self._strict = False  # Use private attribute
        self.invalid_args: list[str] = []  # Store invalid args found during init

    @property
    def strict(self) -> bool:
        """Get the current strict mode setting. Returns True if strict validation of init arguments
        is enabled, False otherwise. Read-only.

        Returns:
            bool: True if strict mode is enabled, False otherwise.

        """
        return self._strict

    @strict.setter
    def strict(self, value: bool) -> None:
        """Set strict mode and validate for invalid arguments if enabled. When True, invalid
        __init__ args raise ValueError. Use at init or before apply.
        """
        if value == self._strict:
            return  # No change needed

        # Only validate if strict is being set to True and we have stored init args
        if value and hasattr(self, "_init_args"):
            # Get the list of valid arguments for this transform
            valid_args = {"p", "strict"}  # Base valid args
            if hasattr(self, "InitSchema"):
                valid_args.update(self.InitSchema.model_fields.keys())

            # Check for invalid arguments
            invalid_args = [name_arg for name_arg in self._init_args if name_arg not in valid_args]

            if invalid_args:
                message = (
                    f"Argument(s) '{', '.join(invalid_args)}' are not valid for transform {self.__class__.__name__}"
                )
                if value:  # In strict mode
                    raise ValueError(message)
                warn(message, stacklevel=2)

        self._strict = value

    def set_random_state(
        self,
        random_generator: np.random.Generator,
        py_random: random.Random,
    ) -> None:
        """Set random state directly from numpy and Python random generators. Used for
        reproducibility and replay. Called by Compose.

        Args:
            random_generator (np.random.Generator): numpy random generator to use
            py_random (random.Random): python random generator to use

        """
        self.random_generator = random_generator
        self.py_random = py_random

    def set_random_seed(self, seed: int | None) -> None:
        """Set random state from a single integer seed. Initializes both numpy and Python random
        generators for reproducibility. Called from __init__.

        Args:
            seed (int | None): Random seed to use

        """
        self.seed = seed
        self.random_generator = np.random.default_rng(seed)
        self.py_random = random.Random(seed)

    def get_dict_with_id(self) -> dict[str, Any]:
        """Return a dictionary representation of the transform with its ID. Used for replay and
        debugging; includes id(self). Same as to_dict plus id.

        Returns:
            dict[str, Any]: Dictionary containing transform parameters and ID.

        """
        d = self.to_dict_private()
        d.update({"id": id(self)})
        return d

    def get_transform_init_args_names(self) -> tuple[str, ...]:
        """Returns names of arguments that are used in __init__ method of the transform.
        Introspects MRO; excludes self and strict.

        This method introspects the entire Method Resolution Order (MRO) to gather the names
        of parameters accepted by the __init__ methods of all parent classes,
        to collect all possible parameters, excluding 'self' and 'strict'
        which are handled separately.
        """
        import inspect

        all_param_names = set()

        for cls in self.__class__.__mro__:
            # Skip the class if it's the base object or doesn't define __init__
            if cls is object or "__init__" not in cls.__dict__:
                continue

            try:
                # Access the class's __init__ method through __dict__ to avoid mypy errors
                init_method = cls.__dict__["__init__"]
                signature = inspect.signature(init_method)
                for name, param in signature.parameters.items():
                    if param.kind in {inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY}:
                        all_param_names.add(name)
            except (ValueError, TypeError):
                continue

        # Exclude 'self' and 'strict'
        return tuple(sorted(all_param_names - {"self", "strict"}))

    def set_processors(self, processors: dict[str, BboxProcessor | KeypointsProcessor]) -> None:
        """Set the processors dictionary used for processing bbox and keypoint transformations.
        Called by Compose when building pipeline.

        Args:
            processors (dict[str, BboxProcessor | KeypointsProcessor]): Dictionary mapping processor
                names to processor instances.

        """
        self.processors = processors

    def get_processor(self, key: str) -> BboxProcessor | KeypointsProcessor | None:
        """Get the processor for a specific key (e.g. bboxes, keypoints). Returns None
        if the key has no processor. Used when applying transforms with params.

        Args:
            key (str): The processor key to retrieve.

        Returns:
            BboxProcessor | KeypointsProcessor | None: The processor instance if found, None otherwise.

        """
        return self.processors.get(key)

    def __call__(self, *args: Any, force_apply: bool = False, **kwargs: Any) -> Any:
        """Apply the transform to the input data. Accepts named kwargs (image, mask, bboxes, etc.);
        returns dict of transformed data.

        Args:
            *args (Any): Positional arguments are not supported and will raise an error.
            force_apply (bool, optional): If True, the transform will be applied regardless of probability.
            **kwargs (Any): Input data to transform as named arguments.

        Returns:
            Any: Transformed data (dict of transformed inputs).

        Raises:
            KeyError: If positional arguments are provided.

        """
        if args:
            msg = "You have to pass data to augmentations as named arguments, for example: aug(image=image)"
            raise KeyError(msg)
        if self.replay_mode:
            if self.applied_in_replay:
                return self.apply_with_params(self.params, **kwargs)
            return kwargs

        # Reset params at the start of each call
        self.params = {}

        if self.should_apply(force_apply=force_apply):
            params = self.get_params()
            params = self.update_transform_params(params=params, data=kwargs)

            if self.targets_as_params:  # check if all required targets are in kwargs.
                missing_keys = set(self.targets_as_params).difference(kwargs.keys())
                if missing_keys and not (missing_keys == {"image"} and "images" in kwargs):
                    msg = f"{self.__class__.__name__} requires {self.targets_as_params} missing keys: {missing_keys}"
                    raise ValueError(msg)

            params_dependent_on_data = self.get_params_dependent_on_data(params=params, data=kwargs)
            params.update(params_dependent_on_data)

            # Store the final params
            self.params = params

            if self.deterministic:
                kwargs[self.save_key][id(self)] = deepcopy(params)
            return self.apply_with_params(params, **kwargs)

        return kwargs

    def get_applied_params(self) -> dict[str, Any]:
        """Returns the parameters that were used in the last transform application; returns empty
        dict if transform was not applied.
        """
        return self.params

    def inverse(self) -> "BasicTransform":
        """Return a new transform that is the mathematical inverse of this one. Useful for TTA to
        revert deterministic transforms. Override in subclasses.

        Useful for TTA (Test-Time Augmentation): apply a deterministic transform to an image
        before inference, then apply its inverse to the predicted mask to bring it back to
        the original image space.

        Only transforms that override `inverse()` support this operation, typically
        group-based transforms with a fixed `group_element` (e.g., D4, RandomRotate90,
        HorizontalFlip, VerticalFlip, Transpose).

        Raises:
            NotImplementedError: If the transform does not support inversion.

        """
        raise NotImplementedError(
            f"{self.__class__.__name__} does not support inverse(). "
            "Only transforms that override `inverse()` can be used for TTA inversion.",
        )

    def should_apply(self, force_apply: bool = False) -> bool:
        """Determine whether to apply the transform based on probability (p) and force_apply flag.
        Used internally before apply_with_params.

        Args:
            force_apply (bool, optional): If True, always apply the transform regardless of probability.

        Returns:
            bool: True if the transform should be applied, False otherwise.

        """
        if self.p <= 0.0:
            return False
        if self.p >= 1.0 or force_apply:
            return True
        return self.py_random.random() < self.p

    def apply_with_params(self, params: dict[str, Any], *args: Any, **kwargs: Any) -> dict[str, Any]:
        """Apply transforms with parameters. Dispatches each target (image, mask, bboxes, etc.) to
        the corresponding apply_* method.
        """
        res = {}
        for key, arg in kwargs.items():
            if key in self._key2func and arg is not None:
                # Handle empty lists for mask-like keys
                if key in {"masks", "masks3d"} and isinstance(arg, (list, tuple)) and not arg:
                    res[key] = arg  # Keep empty list as is
                else:
                    target_function = self._key2func[key]
                    res[key] = target_function(arg, **params)
            else:
                res[key] = arg
        return res

    def set_deterministic(self, flag: bool, save_key: str = "replay") -> "BasicTransform":
        """Set transform to be deterministic. When True, params are saved under save_key for
        replay (e.g. TTA). Returns self for chaining.
        """
        if save_key == "params":
            msg = "params save_key is reserved"
            raise KeyError(msg)

        self.deterministic = flag
        if self.deterministic and self.targets_as_params:
            warn(
                self.get_class_fullname() + " could work incorrectly in ReplayMode for other input data"
                " because its' params depend on targets.",
                stacklevel=2,
            )
        self.save_key = save_key
        return self

    def __repr__(self) -> str:
        state = self.get_base_init_args()
        state.update(self.get_transform_init_args())
        return f"{self.__class__.__name__}({format_args(state)})"

    def apply(self, img: ImageType, *args: Any, **params: Any) -> ImageType:
        """Apply transform on image. Override in subclasses; receives params from get_params and
        get_params_dependent_on_data. Single image only.
        """
        raise NotImplementedError

    @staticmethod
    def _apply_to_batch(
        batch: np.ndarray,
        apply_fn: Callable[[np.ndarray], np.ndarray],
        *,
        ensure_contiguous: bool = False,
    ) -> np.ndarray:
        """Apply a function to each element in a batch with pre-allocation. Uses first element to
        determine output shape; avoids per-call allocation.

        Args:
            batch (np.ndarray): Input batch array of shape (N, ...)
            apply_fn (Callable[[np.ndarray], np.ndarray]): Function to apply to each element
            ensure_contiguous (bool): Whether to ensure C-contiguous output

        Returns:
            np.ndarray: Transformed batch array.

        """
        # Handle empty batch
        if len(batch) == 0:
            return batch

        # Process first element to determine output shape
        first_result = apply_fn(batch[0])

        # Single element case
        if len(batch) == 1:
            result = first_result[np.newaxis]
            return np.require(result, requirements=["C_CONTIGUOUS"]) if ensure_contiguous else result

        # Pre-allocate for remaining elements based on first result
        result_shape = (len(batch), *first_result.shape)
        result = np.empty(result_shape, dtype=first_result.dtype)
        result[0] = first_result

        for i in range(1, len(batch)):
            result[i] = apply_fn(batch[i])

        return np.require(result, requirements=["C_CONTIGUOUS"]) if ensure_contiguous else result

    def apply_to_images(self, images: ImageType, *args: Any, **params: Any) -> ImageType:
        """Apply transform on images. Input shape (N, H, W, C); uses _apply_to_batch with per-image
        apply. Returns same format. Batch API.

        Args:
            images (ImageType): Input images as numpy array of shape:
                - (num_images, height, width, channels)
                - (num_images, height, width) for grayscale
            *args (Any): Additional positional arguments
            **params (Any): Additional parameters specific to the transform

        Returns:
            ImageType: Transformed images as numpy array in the same format as input

        """
        return self._apply_to_batch(images, lambda img: self.apply(img, **params))

    def apply_to_volume(self, volume: VolumeType, *args: Any, **params: Any) -> VolumeType:
        """Apply transform slice by slice to a volume. Delegates to apply_to_images so each slice
        is transformed consistently. Single volume.

        Args:
            volume (VolumeType): Input volume of shape (depth, height, width) or (depth, height, width, channels)
            *args (Any): Additional positional arguments
            **params (Any): Additional parameters specific to the transform

        Returns:
            VolumeType: Transformed volume as numpy array in the same format as input

        """
        return self.apply_to_images(volume, *args, **params)

    def apply_to_volumes(self, volumes: VolumeType, *args: Any, **params: Any) -> VolumeType:
        """Apply transform to multiple volumes. Uses _apply_to_batch; each volume is processed via
        apply_to_volume. Returns same format.
        """
        return self._apply_to_batch(volumes, lambda vol: self.apply_to_volume(vol, *args, **params))

    def get_params(self) -> dict[str, Any]:
        """Returns parameters independent of input data. Override in subclasses to add random
        params (e.g. angle, crop size). Default returns {}.
        """
        return {}

    def update_transform_params(self, params: dict[str, Any], data: dict[str, Any]) -> dict[str, Any]:
        """Updates parameters with input shape and transform-specific params (interpolation, fill,
        fill_mask, bbox_type). Merges get_params output.

        Args:
            params (dict[str, Any]): Parameters to be updated
            data (dict[str, Any]): Input data dictionary containing images/volumes

        Returns:
            dict[str, Any]: Updated parameters dictionary with shape and transform-specific params

        """
        # Extract shape from any available data source
        shape = self._extract_shape_from_data(data)
        if shape is not None:
            params["shape"] = shape

        bbox_processor = self.processors.get("bboxes")
        if isinstance(bbox_processor, BboxProcessor):
            params["bbox_type"] = bbox_processor.params.bbox_type

        # Add transform-specific params
        self._add_transform_specific_params(params)

        return params

    def _extract_shape_from_data(self, data: dict[str, Any]) -> tuple[int, ...] | None:
        """Extract shape information from various data types (image, images, volume, mask, etc.).
        Used for params that need H, W. Returns tuple or None.
        """
        # Check images/volumes first (most common)
        if "image" in data:
            return data["image"].shape
        if "images" in data:
            return data["images"][0].shape
        if "volume" in data:
            # For volumes, take a slice to get 2D shape
            return data["volume"][0].shape  # Take first slice (works for both DHWC and DHW)
        if "volumes" in data:
            # For batch of volumes
            return data["volumes"][0][0].shape  # Take first slice of first volume (works for both NDHWC and NDHW)
        if "mask" in data:
            return data["mask"].shape
        if "masks" in data:
            return data["masks"][0].shape
        if "mask3d" in data:
            # For 3D masks, take a slice to get 2D shape
            return data["mask3d"][0].shape  # Take first slice (works for both DHWC and DHW)
        if "masks3d" in data:
            # For batch of 3D masks
            return data["masks3d"][0][0].shape  # Take first slice of first mask (works for both NDHWC and NDHW)
        return None

    def _add_transform_specific_params(self, params: dict[str, Any]) -> None:
        """Add transform-specific parameters to params dict (interpolation, fill, fill_mask).
        Called from update_transform_params. Mutates params in place.
        """
        if hasattr(self, "interpolation"):
            params["interpolation"] = self.interpolation
        if hasattr(self, "fill"):
            params["fill"] = self.fill
        if hasattr(self, "fill_mask"):
            params["fill_mask"] = self.fill_mask

    def get_params_dependent_on_data(self, params: dict[str, Any], data: dict[str, Any]) -> dict[str, Any]:
        """Returns parameters dependent on input data (e.g. crop coordinates from image shape).
        Override in subclasses; default returns params unchanged.
        """
        return params

    @property
    def targets(self) -> dict[str, Callable[..., Any]]:
        """Get mapping of target keys to their corresponding processing functions (e.g. image ->
        apply, mask -> apply_to_mask). Subclasses override.

        Returns:
            dict[str, Callable[..., Any]]: Dictionary mapping target keys to their processing functions.

        """
        # mapping for targets and methods for which they depend
        # for example:
        # >>  {"image": self.apply}
        # >>  {"masks": self.apply_to_masks}
        raise NotImplementedError

    def apply_to_user_data(self, data: Any, **params: Any) -> Any:
        """Apply transform to user-defined data. By default returns data unchanged (passthrough).
        Override to update custom keys (e.g. captions) from params.

        By default, returns the data unchanged (passthrough). Override in a subclass to
        update arbitrary user data in response to geometric or photometric transforms.

        Args:
            data (Any): Arbitrary user-defined data of any type.
            **params (Any): Transform parameters (same as passed to other apply_* methods).

        Returns:
            Any: The (optionally modified) user data. Must return the same type as the input.

        Examples:
            >>> import albumentations as A
            >>> class FlipAwareTransform(A.HorizontalFlip):
            ...     def apply_to_user_data(self, data: dict, **params) -> dict:
            ...         return {"caption": data["caption"].replace("left", "right")}

        """
        return data

    def _set_keys(self) -> None:
        """Set _available_keys and _key2func from _targets and targets. Adds user_data as
        passthrough. Called from __init__. Override targets in subclass.
        """
        if not hasattr(self, "_targets"):
            self._available_keys = set()
        else:
            self._available_keys = {
                target.value.lower()
                for target in (self._targets if isinstance(self._targets, tuple) else [self._targets])
            }
        self._available_keys.update(self.targets.keys())
        self._key2func = {key: self.targets[key] for key in self._available_keys if key in self.targets}
        # user_data is always available regardless of _targets - passthrough by default
        self._available_keys.add("user_data")
        self._key2func["user_data"] = self.apply_to_user_data

    @property
    def available_keys(self) -> set[str]:
        """Returns set of available keys (target names this transform can process). Includes
        built-in targets and add_targets additions.
        """
        return self._available_keys

    def add_targets(self, additional_targets: dict[str, str]) -> None:
        """Register additional targets transformed like an existing one (e.g. {'image2': 'image'}).
        Need at least 'image' in pipeline.

        Args:
            additional_targets (dict[str, str]): keys - new target name, values
                - old target name. ex: {'image2': 'image'}

        """
        for k, v in additional_targets.items():
            if k in self._additional_targets and v != self._additional_targets[k]:
                raise ValueError(
                    f"Trying to overwrite existed additional targets. "
                    f"Key={k} Exists={self._additional_targets[k]} New value: {v}",
                )
            if v in self._available_keys:
                self._additional_targets[k] = v
                self._key2func[k] = self._key2func[v]
                self._available_keys.add(k)

    @property
    def targets_as_params(self) -> list[str]:
        """Targets used to get params dependent on targets. Used to check input has all required
        targets before apply. Override to list keys (e.g. ['image']).
        """
        return []

    @classmethod
    def get_class_fullname(cls) -> str:
        """Get the full qualified name of the class. Returns shortest fullname for serialization
        (e.g. albumentations.HorizontalFlip).

        Returns:
            str: The shortest class fullname.

        """
        return get_shortest_class_fullname(cls)

    @classmethod
    def is_serializable(cls) -> bool:
        """Check if the transform class is serializable. True for all registered transforms; used
        by serialization to skip non-serializable classes.

        Returns:
            bool: True if the class is serializable, False otherwise.

        """
        return True

    def get_base_init_args(self) -> dict[str, Any]:
        """Returns base init args (e.g. p) for serialization. Subclasses may override
        to add more; merged into to_dict_private output.
        """
        return {"p": self.p}

    def get_transform_init_args(self) -> dict[str, Any]:
        """Get transform initialization arguments for serialization. Returns dict of init param
        names and values, excluding empty containers and seed.

        Returns a dictionary of parameter names and their values, excluding parameters
        that are not actually set on the instance or that shouldn't be serialized.
        """
        # Get the parameter names
        arg_names = self.get_transform_init_args_names()

        # Create a dictionary of parameter values
        args = {}
        for name in arg_names:
            # Only include parameters that are actually set as instance attributes
            # and have non-default values
            if hasattr(self, name):
                value = getattr(self, name)
                # Skip attributes that are basic containers with no content
                if not (isinstance(value, (list, dict, tuple, set)) and len(value) == 0):
                    args[name] = value

        # Remove seed explicitly (it's not meant to be serialized)
        args.pop("seed", None)

        return args

    def to_dict_private(self) -> dict[str, Any]:
        """Returns a dictionary representation of the transform for serialization.
        Excludes internal parameters; includes __class_fullname__ and init args.
        """
        state = {"__class_fullname__": self.get_class_fullname()}
        state.update(self.get_base_init_args())

        # Get transform init args (our improved method handles all types of transforms)
        transform_args = self.get_transform_init_args()

        # Add transform args to state
        state.update(transform_args)

        # Remove strict from serialization
        state.pop("strict", None)

        return state


class DualTransform(BasicTransform):
    """Base class for transforms that apply to both image and annotations (masks, bboxes, keypoints),
    keeping them spatially consistent.

    When a transform is applied to an image, all associated entities (masks, bounding boxes, keypoints) are
    such as masks, bounding boxes, and keypoints. This class ensures that when a transform is applied to an image,
    all associated entities are transformed accordingly to maintain consistency between the image and its annotations.

    Class Attributes:
        _supported_bbox_types (set[str]): Set of supported bounding box types.
            Valid values: {"hbb"} for axis-aligned boxes only, {"hbb", "obb"} for both axis-aligned
            and oriented boxes. Default: {"hbb"}. Transforms that support OBB should override this.

    Methods:
        apply(img: np.ndarray, **params: Any) -> np.ndarray:
            Apply the transform to the image.

            img: Input image of shape (H, W, C).
            **params: Additional parameters specific to the transform.

            Returns Transformed image of the same shape as input.

        apply_to_images(images: ImageType, **params: Any) -> ImageType:
            Apply the transform to multiple images.

            images: Input images of shape (N, H, W, C).
            **params: Additional parameters specific to the transform.

            Returns Transformed images in the same format as input.

        apply_to_mask(mask: ImageType, **params: Any) -> ImageType:
            Apply the transform to a mask.

            mask: Input mask of shape (H, W), (H, W, C) for multi-channel masks
            **params: Additional parameters specific to the transform.

            Returns Transformed mask in the same format as input.

        apply_to_masks(masks: ImageType, **params: Any) -> ImageType:
            Apply the transform to multiple masks.

            masks: Array of shape (N, H, W) or (N, H, W, C) where N is number of masks
            **params: Additional parameters specific to the transform.
            Returns Transformed masks in the same format as input.

        apply_to_keypoints(keypoints: np.ndarray, **params: Any) -> np.ndarray:
            Apply the transform to keypoints.

            keypoints: Array of shape (N, 2+) where N is the number of keypoints.
                **params: Additional parameters specific to the transform.
            Returns Transformed keypoints array of shape (N, 2+).

        apply_to_bboxes(bboxes: np.ndarray, **params: Any) -> np.ndarray:
            Apply the transform to bounding boxes.

            bboxes: Array of shape (N, 4+) where N is the number of bounding boxes,
                    and each row is in the format [x_min, y_min, x_max, y_max].
            **params: Additional parameters specific to the transform.

            Returns Transformed bounding boxes array of shape (N, 4+).

        apply_to_volume(volume: VolumeType, **params: Any) -> VolumeType:
            Apply the transform to a volume.

            volume: Input volume of shape (D, H, W, C).
            **params: Additional parameters specific to the transform.

            Returns Transformed volume of the same shape as input.

        apply_to_volumes(volumes: VolumeType, **params: Any) -> VolumeType:
            Apply the transform to multiple volumes.

            volumes: Input volumes of shape (N, D, H, W, C).
            **params: Additional parameters specific to the transform.

            Returns Transformed volumes in the same format as input.

        apply_to_mask3d(mask: VolumeType, **params: Any) -> VolumeType:
            Apply the transform to a 3D mask.

            mask: Input 3D mask of shape (D, H, W) or (D, H, W, C)
            **params: Additional parameters specific to the transform.

            Returns Transformed 3D mask in the same format as input.

        apply_to_masks3d(masks: VolumeType, **params: Any) -> VolumeType:
            Apply the transform to multiple 3D masks.

            masks: Input 3D masks of shape (N, D, H, W) or (N, D, H, W, C)
            **params: Additional parameters specific to the transform.

            Returns Transformed 3D masks in the same format as input.

    Note:
        - All `apply_*` methods should maintain the input shape and format of the data.
        - When applying transforms to masks, ensure that discrete values (e.g., class labels) are preserved.
        - For keypoints and bounding boxes, the transformation should maintain their relative positions
            with respect to the transformed image.
        - The difference between `apply_to_mask` and `apply_to_masks` is mainly in how they handle 3D arrays:
            `apply_to_mask` treats a 3D array as a multi-channel mask, while `apply_to_masks` treats it as
            multiple single-channel masks.

    """

    _supported_bbox_types: frozenset[str] = frozenset({"hbb"})  # Default: only axis-aligned boxes

    @property
    def targets(self) -> dict[str, Callable[..., Any]]:
        """Get mapping of target keys to their corresponding processing functions for DualTransform
        (image, mask, bboxes, keypoints, etc.).

        Returns:
            dict[str, Callable[..., Any]]: Dictionary mapping target keys to their processing functions.

        """
        # Note: keypoint label swapping is handled within apply_to_keypoints
        # No separate targets needed for label fields
        return {
            "image": self.apply,
            "images": self.apply_to_images,
            "mask": self.apply_to_mask,
            "masks": self.apply_to_masks,
            "mask3d": self.apply_to_mask3d,
            "masks3d": self.apply_to_masks3d,
            "bboxes": self.apply_to_bboxes,
            "keypoints": self.apply_to_keypoints,
            "volume": self.apply_to_images,
            "volumes": self.apply_to_volumes,
            "user_data": self.apply_to_user_data,
        }

    def apply_to_keypoints(self, keypoints: np.ndarray, *args: Any, **params: Any) -> np.ndarray:
        msg = f"Method apply_to_keypoints is not implemented in class {self.__class__.__name__}"
        raise NotImplementedError(msg)

    def apply_to_bboxes(self, bboxes: np.ndarray, *args: Any, **params: Any) -> np.ndarray:
        raise NotImplementedError(f"BBoxes not implemented for {self.__class__.__name__}")

    def apply_to_mask(self, mask: ImageType, *args: Any, **params: Any) -> ImageType:
        return self.apply(mask, *args, **params)

    def apply_to_masks(self, masks: ImageType, *args: Any, **params: Any) -> ImageType:
        if masks.size == 0:
            return masks
        return self._apply_to_batch(masks, lambda mask: self.apply_to_mask(mask, *args, **params))

    @batch_transform("spatial")
    def apply_to_mask3d(self, mask3d: VolumeType, *args: Any, **params: Any) -> VolumeType:
        return self.apply_to_mask(mask3d, *args, **params)

    @batch_transform("spatial")
    def apply_to_masks3d(self, masks3d: VolumeType, *args: Any, **params: Any) -> VolumeType:
        return self._apply_to_batch(masks3d, lambda mask3d: self.apply_to_mask3d(mask3d, **params))

    def _get_label_transform_name(self, **params: Any) -> str | None:
        """Get the transform name to use for label mapping. For most transforms returns class
        name; for D4/SquareSymmetry maps group_element to base name.

        For most transforms, this is just the class name. For D4/SquareSymmetry,
        we map the group element to the corresponding base transform name.

        Args:
            **params (Any): Transform parameters, may contain group_element for D4 transforms

        Returns:
            str | None: Transform name to use for label mapping, or None if no mapping should be applied

        """
        class_name = self.__class__.__name__

        # Handle D4 and SquareSymmetry transforms (including subclasses)
        if class_name in ("D4", "SquareSymmetry") or any(
            base.__name__ in ("D4", "SquareSymmetry") for base in self.__class__.__mro__
        ):
            group_element = params.get("group_element", "e")
            # Map D4 group elements to base transform names
            d4_to_base_transform = {
                "h": "HorizontalFlip",
                "v": "VerticalFlip",
                "t": "Transpose",
                "hvt": "Transpose",  # Anti-diagonal is also a transpose-like operation
                "e": None,  # Identity - no label swapping
                "r90": None,  # Rotations don't change semantic labels
                "r180": None,
                "r270": None,
            }
            mapped_name = d4_to_base_transform.get(group_element)
            return mapped_name or class_name

        # Only parity-changing transforms should apply label mappings
        parity_changing_transforms = {"HorizontalFlip", "VerticalFlip", "Transpose"}
        return class_name if class_name in parity_changing_transforms else None

    def _apply_label_mapping_to_keypoints(self, keypoints: np.ndarray, **params: Any) -> np.ndarray:
        """Apply label mapping by reordering entire keypoint rows. For keypoint regression, row
        index encodes semantics; flip/transpose swap rows via mapping.

        For keypoint regression tasks, the row index encodes semantic meaning
        (e.g., row 0 = left eye heatmap). On transforms like HorizontalFlip,
        we need to swap entire rows, not just relabel them.

        Args:
            keypoints (np.ndarray): Keypoints array with potential label columns attached
            **params (Any): Transform parameters

        Returns:
            np.ndarray: Keypoints array with rows reordered based on label mapping

        """
        # Get the keypoint processor
        processor = self.processors.get("keypoints") if hasattr(self, "processors") else None
        if not processor or not hasattr(processor, "encoded_label_mappings"):
            return keypoints

        # Check if there are label fields and the array has extra columns
        if not processor.params.label_fields or keypoints.size == 0 or keypoints.shape[1] <= 5:
            return keypoints

        transform_name = self._get_label_transform_name(**params)
        if transform_name is None or transform_name not in processor.encoded_label_mappings:
            return keypoints

        # Only copy if we actually have mappings to apply
        field_mappings = processor.encoded_label_mappings[transform_name]
        if not field_mappings:
            return keypoints

        return self._swap_keypoint_rows_by_labels(keypoints, processor.params.label_fields, field_mappings)

    def _swap_keypoint_rows_by_labels(
        self,
        keypoints: np.ndarray,
        label_fields: Sequence[str],
        field_mappings: dict[str, dict[int, int]],
    ) -> np.ndarray:
        """Swap keypoint rows based on label mappings. Used when transform changes left/right or
        similar; swaps entire rows so coords and labels stay consistent.

        Args:
            keypoints (np.ndarray): Keypoints array with label columns
            label_fields (Sequence[str]): List of label field names
            field_mappings (dict[str, dict[int, int]]): Mapping of field names to label swaps

        Returns:
            np.ndarray: Keypoints array with rows swapped

        """
        result = keypoints.copy()
        label_col_start = 5  # After [x, y, z, angle, scale]

        # For each label field with mapping, perform row swapping
        for i, label_field in enumerate(label_fields):
            if label_field in field_mappings:
                col_idx = label_col_start + i
                if col_idx < keypoints.shape[1]:
                    mapping = field_mappings[label_field]
                    if mapping:  # Only process if mapping is not empty
                        result = self._apply_single_field_mapping(result, col_idx, mapping)
                        # Only apply mapping for the first label field that has mappings
                        break

        return result

    def _apply_single_field_mapping(
        self,
        keypoints: np.ndarray,
        col_idx: int,
        mapping: dict[int, int],
    ) -> np.ndarray:
        """Apply label mapping to a single label column. Swaps rows for paired labels or updates
        unpaired; used internally by _swap_keypoint_rows_by_labels.

        Args:
            keypoints (np.ndarray): Keypoints array
            col_idx (int): Column index of the label field
            mapping (dict[int, int]): Label swap mapping

        Returns:
            np.ndarray: Keypoints array with rows swapped

        """
        col_data = keypoints[:, col_idx].astype(int)
        processed_labels = set()

        for from_label, to_label in mapping.items():
            if from_label in processed_labels or to_label in processed_labels:
                continue

            from_indices = np.where(col_data == from_label)[0]
            to_indices = np.where(col_data == to_label)[0]

            # If both labels exist in data, swap entire rows
            if len(from_indices) > 0 and len(to_indices) > 0:
                # Swap entire rows (coordinates + all labels)
                temp_rows = keypoints[from_indices].copy()
                keypoints[from_indices] = keypoints[to_indices]
                keypoints[to_indices] = temp_rows
                processed_labels.add(from_label)
                processed_labels.add(to_label)
            # If only from_label exists (unpaired), just update its label
            elif len(from_indices) > 0:
                keypoints[from_indices, col_idx] = to_label
                processed_labels.add(from_label)

        return keypoints

    def apply_with_params(self, params: dict[str, Any], *args: Any, **kwargs: Any) -> dict[str, Any]:
        """Apply transforms with parameters, including automatic keypoint label swapping. After
        super().apply_with_params, applies _apply_label_mapping_to_keypoints.
        """
        res = super().apply_with_params(params, *args, **kwargs)

        # Apply label mapping to keypoints if they were transformed
        if "keypoints" in res and res["keypoints"] is not None:
            res["keypoints"] = self._apply_label_mapping_to_keypoints(res["keypoints"], **params)

        return res


class ImageOnlyTransform(BasicTransform):
    """Transform applied to image (and volume) only. Does not transform masks, bboxes, or
    keypoints; use DualTransform for those.
    """

    _targets = (Targets.IMAGE, Targets.VOLUME)

    @property
    def targets(self) -> dict[str, Callable[..., Any]]:
        """Get mapping of target keys to their corresponding processing functions for
        ImageOnlyTransform (image, images, volume, volumes, user_data).

        Returns:
            dict[str, Callable[..., Any]]: Dictionary mapping target keys to their processing functions.

        """
        return {
            "image": self.apply,
            "images": self.apply_to_images,
            "volume": self.apply_to_volume,
            "volumes": self.apply_to_volumes,
            "user_data": self.apply_to_user_data,
        }


class NoOp(DualTransform):
    """Identity transform (does nothing). Passes all targets through unchanged. Use as placeholder
    or in conditional pipelines.

    Targets:
        image, mask, bboxes, keypoints, volume, mask3d

    Image types:
        uint8, float32

    Supported bboxes:
        hbb, obb

    Examples:
        >>> import numpy as np
        >>> import albumentations as A
        >>>
        >>> # Prepare sample data
        >>> image = np.random.randint(0, 256, (100, 100, 3), dtype=np.uint8)
        >>> mask = np.random.randint(0, 2, (100, 100), dtype=np.uint8)
        >>> bboxes = np.array([[10, 10, 50, 50], [40, 40, 80, 80]], dtype=np.float32)
        >>> bbox_labels = [1, 2]
        >>> keypoints = np.array([[20, 30], [60, 70]], dtype=np.float32)
        >>> keypoint_labels = [0, 1]
        >>>
        >>> # Create transform pipeline with NoOp
        >>> transform = A.Compose([
        ...     A.NoOp(p=1.0),  # Always applied, but does nothing
        ... ], bbox_params=A.BboxParams(coord_format='pascal_voc', label_fields=['bbox_labels']),
        ...    keypoint_params=A.KeypointParams(coord_format='xy', label_fields=['keypoint_labels']))
        >>>
        >>> # Apply the transform
        >>> transformed = transform(
        ...     image=image,
        ...     mask=mask,
        ...     bboxes=bboxes,
        ...     bbox_labels=bbox_labels,
        ...     keypoints=keypoints,
        ...     keypoint_labels=keypoint_labels
        ... )
        >>>
        >>> # Verify nothing has changed
        >>> np.array_equal(image, transformed['image'])  # True
        >>> np.array_equal(mask, transformed['mask'])  # True
        >>> np.array_equal(bboxes, transformed['bboxes'])  # True
        >>> np.array_equal(keypoints, transformed['keypoints'])  # True
        >>> bbox_labels == transformed['bbox_labels']  # True
        >>> keypoint_labels == transformed['keypoint_labels']  # True
        >>>
        >>> # NoOp is often used as a placeholder or for testing
        >>> # For example, in conditional transforms:
        >>> condition = False  # Some condition
        >>> transform = A.Compose([
        ...     A.HorizontalFlip(p=1.0) if condition else A.NoOp(p=1.0)
        ... ])

    """

    _targets = ALL_TARGETS
    _supported_bbox_types: frozenset[str] = frozenset({"hbb", "obb"})  # NoOp passes all bbox types

    def apply_to_keypoints(self, keypoints: np.ndarray, **params: Any) -> np.ndarray:
        return keypoints

    def apply_to_bboxes(self, bboxes: np.ndarray, **params: Any) -> np.ndarray:
        return bboxes

    def apply(self, img: ImageType, **params: Any) -> ImageType:
        return img

    def apply_to_mask(self, mask: ImageType, **params: Any) -> ImageType:
        return mask

    def apply_to_volume(self, volume: VolumeType, **params: Any) -> VolumeType:
        return volume

    def apply_to_volumes(self, volumes: VolumeType, **params: Any) -> VolumeType:
        return volumes

    def apply_to_mask3d(self, mask3d: VolumeType, **params: Any) -> VolumeType:
        return mask3d

    def apply_to_masks3d(self, masks3d: VolumeType, **params: Any) -> VolumeType:
        return masks3d


class Transform3D(DualTransform):
    """Base class for all 3D transforms. Inherits from DualTransform; applies to volumes,
    masks3d, keypoints. Override apply_to_volume and apply_to_mask3d.

    Transform3D inherits from DualTransform because 3D transforms can be applied to both
    volumes and masks, similar to how 2D DualTransforms work with images and masks.

    Targets:
        volume: 3D numpy array of shape (D, H, W, C)
        volumes: Batch of 3D arrays of shape (N, D, H, W, C)
        mask: 3D numpy array of shape (D, H, W) or (D, H, W, C)
        masks: Batch of 3D arrays of shape (N, D, H, W) or (N, D, H, W, C)
        keypoints: 3D numpy array of shape (N, 3)
    """

    def apply_to_volume(self, volume: VolumeType, *args: Any, **params: Any) -> VolumeType:
        """Apply transform to single 3D volume. Override in subclasses; input shape (D, H, W, C)
        or (D, H, W). Returns same shape and dtype.
        """
        raise NotImplementedError

    @batch_transform("spatial", keep_depth_dim=True)
    def apply_to_volumes(self, volumes: VolumeType, *args: Any, **params: Any) -> VolumeType:
        """Apply transform to batch of 3D volumes. Uses batch_transform with keep_depth_dim;
        each volume passed to apply_to_volume.
        """
        return self.apply_to_volume(volumes, *args, **params)

    def apply_to_mask3d(self, mask3d: VolumeType, *args: Any, **params: Any) -> VolumeType:
        """Apply transform to a single 3D mask. Delegates to apply_to_volume. Input shape (D, H, W) or
        (D, H, W, C). Output shape unchanged. For VolumeTransform.
        """
        return self.apply_to_volume(mask3d, *args, **params)

    @batch_transform("spatial", keep_depth_dim=True)
    def apply_to_masks3d(self, masks3d: VolumeType, *args: Any, **params: Any) -> VolumeType:
        """Apply transform to batch of 3D masks. Uses batch_transform with keep_depth_dim;
        each mask passed to apply_to_mask3d. Same shape.
        """
        return self.apply_to_mask3d(masks3d, *args, **params)

    @property
    def targets(self) -> dict[str, Callable[..., Any]]:
        return {
            "volume": self.apply_to_volume,
            "volumes": self.apply_to_volumes,
            "mask3d": self.apply_to_mask3d,
            "masks3d": self.apply_to_masks3d,
            "keypoints": self.apply_to_keypoints,
            "user_data": self.apply_to_user_data,
        }


class CustomTransformsApplyMixin:
    """Mixin that auto-registers custom apply_to_<X> methods as handlers for data key <X>.
    Place before base in MRO so _set_keys discovers them.

    Define methods named `apply_to_<key>` in your transform subclass; they are
    discovered at init time and routed through the standard `apply_with_params`
    pipeline. Custom targets receive the same params from `get_params`, respect
    the `p=` probability, and compose correctly with Compose and ReplayCompose.

    Placement in inheritance list
        Must come BEFORE the albumentations base class so MRO resolves
        `_set_keys` first::

            class MyTransform(CustomTransformsApplyMixin, A.DualTransform):
                def apply_to_label(self, label, **params):
                    return (label + params["factor']) % 4

    Registration rules
        Methods named `apply_to_<X>` are registered if they are:
        - Defined in the concrete subclass or any class between it and this mixin in the MRO
        - Not already covered by `self.targets` (built-ins take priority)
        - Not `apply_to_user_data` (handled separately by the base)

    Examples:
        >>> import numpy as np
        >>> import albumentations as A
        >>> image = np.random.randint(0, 256, (64, 64, 3), dtype=np.uint8)
        >>> mask = np.random.randint(0, 2, (64, 64), dtype=np.uint8)
        >>>
        >>> class Rotate90WithLabel(A.CustomTransformsApplyMixin, A.DualTransform):
        ...     def get_params(self):
        ...         return {"k": 1}
        ...     def apply(self, img, k=0, **p):
        ...         return np.rot90(img, k)
        ...     def apply_to_mask(self, mask, k=0, **p):
        ...         return np.rot90(mask, k)
        ...     def apply_to_label(self, label, k=0, **p):
        ...         return (label + k) % 4
        >>>
        >>> transform = A.Compose([Rotate90WithLabel(p=1.0)])
        >>> out = transform(image=image, mask=mask, label=0)
        >>> out["label"]
        1

    """

    _APPLY_PREFIX = "apply_to_"
    _EXCLUDED_KEYS = frozenset({"user_data"})
    _key2func: dict[str, Any]
    _available_keys: set[str]

    def _set_keys(self) -> None:
        # Build _key2func from self.targets using base class
        super()._set_keys()  # type: ignore[misc]

        # Search apply_to_<X> functions defined within the child class
        for name, method in inspect.getmembers(self, predicate=inspect.ismethod):
            if not name.startswith(self._APPLY_PREFIX):
                continue
            key = name[len(self._APPLY_PREFIX) :]
            if key in self._EXCLUDED_KEYS:
                continue
            if key in self._key2func:  # built-in already registered
                continue
            if not self._is_user_defined(name):
                continue
            self._available_keys.add(key)
            self._key2func[key] = method

    def _is_user_defined(self, method_name: str) -> bool:
        """True if method_name is defined on subclass or parents before mixin in MRO (not from albumentations base).
        Used to register only user-defined apply_to_<X>.
        """
        for mro_class in type(self).__mro__:
            if mro_class is CustomTransformsApplyMixin:
                break
            if method_name in mro_class.__dict__:
                return True
        return False
