ng_data_config

class Mode(str, BaseEnum):
    """Dataset mode."""

    TRAINING = "training"
    VALIDATING = "validating"
    PREDICTING = "predicting"

class NGDataConfig(BaseModel):
    """Next-Generation Dataset configuration.

    NGDataConfig are used for both training and prediction, with the patching strategy
    determining how the data is processed. Note that `random` is the only patching
    strategy compatible with training, while `tiled` and `whole` are only used for
    prediction.

    All supported transforms are defined in the SupportedTransform enum.
    """

    # Pydantic class configuration
    model_config = ConfigDict(
        validate_assignment=True,
    )

    # Dataset configuration
    mode: Mode
    """Dataset mode, either training, validating or predicting."""

    data_type: Literal["array", "tiff", "zarr", "czi", "custom"]
    """Type of input data."""

    axes: str
    """Axes of the data, as defined in SupportedAxes."""

    # TODO: update docs for stratified patching
    patching: PatchingConfig = Field(..., discriminator="name")
    """Patching strategy to use. Note that `random` is the only supported strategy for
    training, while `tiled` and `whole` are only used for prediction."""

    normalization: NormalizationConfig = Field(...)
    """Normalization configuration to use."""

    # Optional fields
    batch_size: int = Field(default=1, ge=1, validate_default=True)
    """Batch size for training."""

    in_memory: bool = Field(default_factory=default_in_memory, validate_default=True)
    """Whether to load all data into memory. This is only supported for 'array',
    'tiff' and 'custom' data types. Must be `True` for `array`. If `None`, defaults to
    `True` for 'array', 'tiff' and `custom`, and `False` for 'zarr' and 'czi' data
    types."""

    n_val_patches: int = Field(default=8, ge=0, validate_default=True)
    """The number of patches to set aside for validation during training. This parameter
    will be ignored if separate validation data is specified for training."""

    channels: Sequence[int] | None = Field(default=None)
    """Channels to use from the data. If `None`, all channels are used."""

    patch_filter: PatchFilterConfig | None = Field(default=None, discriminator="name")
    """Patch filter to apply when using random patching. Only available if
    mode is `training`."""

    coord_filter: CoordFilterConfig | None = Field(default=None, discriminator="name")
    """Coordinate filter to apply when using random patching. Only available if
    mode is `training`."""

    patch_filter_patience: int = Field(default=5, ge=1)
    """Number of consecutive patches not passing the filter before accepting the next
    patch."""

    augmentations: Sequence[Union[XYFlipConfig, XYRandomRotate90Config]] = Field(
        default=(
            XYFlipConfig(),
            XYRandomRotate90Config(),
        ),
        validate_default=True,
    )
    """List of augmentations to apply to the data, available transforms are defined
    in SupportedTransform."""

    train_dataloader_params: dict[str, Any] = Field(
        default={"shuffle": True}, validate_default=True
    )
    """Dictionary of PyTorch training dataloader parameters. The dataloader parameters,
    should include the `shuffle` key, which is set to `True` by default. We strongly
    recommend to keep it as `True` to ensure the best training results."""

    val_dataloader_params: dict[str, Any] = Field(default={})
    """Dictionary of PyTorch validation dataloader parameters."""

    pred_dataloader_params: dict[str, Any] = Field(default={})
    """Dictionary of PyTorch prediction dataloader parameters."""

    seed: int = Field(default_factory=generate_random_seed, gt=0)
    """Random seed for reproducibility. If not specified, a random seed is generated."""

    @field_validator("axes")
    @classmethod
    def axes_valid(cls, axes: str, info: ValidationInfo) -> str:
        """
        Validate axes.

        Axes must:
        - be a combination of 'STCZYX'
        - not contain duplicates
        - contain at least 2 contiguous axes: X and Y
        - contain at most 4 axes
        - not contain both S and T axes

        Parameters
        ----------
        axes : str
            Axes to validate.
        info : ValidationInfo
            Validation information.

        Returns
        -------
        str
            Validated axes.

        Raises
        ------
        ValueError
            If axes are not valid.
        """
        if "data_type" not in info.data:
            raise ValueError(
                "Validation for `data_type` may have failed. Check for typos or "
                "missing field."
            )

        # Additional validation for CZI files
        if info.data["data_type"] == "czi":
            if not check_czi_axes_validity(axes):
                raise ValueError(
                    f"Invalid axes '{axes}'. Axes must be in the "
                    f"`SC(Z/T)YX` format, where Z or T are optional, and S and C can be"
                    f" singleton dimensions, but must be provided."
                )
        else:
            check_axes_validity(axes)

        return axes

    @field_validator("in_memory")
    @classmethod
    def validate_in_memory_with_data_type(cls, in_memory: bool, info: Any) -> bool:
        """
        Validate that in_memory is compatible with data_type.

        `in_memory` can only be True for 'array', 'tiff' and 'custom' data types.

        Parameters
        ----------
        in_memory : bool
            Whether to load data into memory.
        info : Any
            Additional information about the field being validated.

        Returns
        -------
        bool
            Validated in_memory value.

        Raises
        ------
        ValueError
            If in_memory is True for unsupported data types.
        """
        data_type = info.data.get("data_type")

        if in_memory and data_type in ("czi", "zarr"):
            raise ValueError(f"`in_memory` not supported for `data_type` {data_type}.")

        if not in_memory and data_type == "array":
            raise ValueError('`in_memory` must be True for "array" `data_type`.')

        return in_memory

    @field_validator("channels", mode="before")
    @classmethod
    def validate_channels(
        cls,
        channels: Sequence[int] | None,
        info: ValidationInfo,
    ) -> Sequence[int] | None:
        """
        Validate channels.

        Channels must be a sequence of non-negative integers without duplicates. If
        channels are not `None`, then `C` must be present in the axes.

        Parameters
        ----------
        channels : Sequence of int or None
            Channels to validate.
        info : ValidationInfo
            Validation information.

        Returns
        -------
        Sequence of int or None
            Validated channels.

        Raises
        ------
        ValueError
            If channels are not valid.
        """
        if channels is not None:
            if "C" not in info.data["axes"]:
                raise ValueError(
                    "Channels must be `None` if 'C' is not present in `axes`."
                )

            if isinstance(channels, int):
                channels = [channels]

            if not isinstance(channels, Sequence):
                raise ValueError("Channels must be a sequence of integers.")

            if len(channels) == 0:
                return None

            if not all(isinstance(ch, int) for ch in channels):
                raise ValueError("Channels must be integers.")

            if any(ch < 0 for ch in channels):
                raise ValueError("Channels must be non-negative integers.")

            if len(set(channels)) != len(channels):
                raise ValueError("Channels must not contain duplicates.")
        return channels

    @field_validator("patching")
    @classmethod
    def validate_patching_strategy_against_mode(
        cls, patching: PatchingConfig, info: ValidationInfo
    ) -> PatchingConfig:
        """
        Validate that the patching strategy is compatible with the dataset mode.

        - If mode is `training`, patching strategy must be `random` or `stratified`.
        - If mode is `validating`, patching must be `fixed_random`.
        - If mode is `predicting`, patching strategy must be `tiled` or `whole`.

        Parameters
        ----------
        patching : PatchingStrategies
            Patching strategy to validate.
        info : ValidationInfo
            Validation information.

        Returns
        -------
        PatchingStrategies
            Validated patching strategy.

        Raises
        ------
        ValueError
            If the patching strategy is not compatible with the dataset mode.
        """
        mode = info.data["mode"]
        if mode == Mode.TRAINING:
            if patching.name not in ["random", "stratified"]:
                raise ValueError(
                    f"Patching strategy '{patching.name}' is not compatible with "
                    f"mode '{mode.value}'. Use 'stratified' or 'random' for training."
                )
        elif mode == Mode.VALIDATING:
            if patching.name != "fixed_random":
                raise ValueError(
                    f"Patching strategy '{patching.name}' is not compatible with "
                    f"mode '{mode.value}'. Use 'fixed_random' for validating."
                )
        elif mode == Mode.PREDICTING:
            if patching.name not in ["tiled", "whole"]:
                raise ValueError(
                    f"Patching strategy '{patching.name}' is not compatible with "
                    f"mode '{mode.value}'. Use 'tiled' or 'whole' for predicting."
                )
        return patching

    @field_validator("patch_filter", "coord_filter")
    @classmethod
    def validate_filters_against_mode(
        cls,
        filter_obj: PatchFilterConfig | CoordFilterConfig | None,
        info: ValidationInfo,
    ) -> PatchFilterConfig | CoordFilterConfig | None:
        """
        Validate that the filters are only used during training.

        Parameters
        ----------
        filter_obj : PatchFilters or CoordFilters or None
            Filter to validate.
        info : ValidationInfo
            Validation information.

        Returns
        -------
        PatchFilters or CoordFilters or None
            Validated filter.

        Raises
        ------
        ValueError
            If a filter is used in a mode other than training.
        """
        mode = info.data["mode"]
        if filter_obj is not None and mode != Mode.TRAINING:
            raise ValueError(
                f"Filtering '{filter_obj.name}' only allowed in 'training' mode, "
                f"got mode '{mode.value}'."
            )
        return filter_obj

    @field_validator(
        "train_dataloader_params",
        "val_dataloader_params",
        "pred_dataloader_params",
        mode="after",
    )
    @classmethod
    def batch_size_not_in_dataloader_params(
        cls, dataloader_params: dict[str, Any]
    ) -> dict[str, Any]:
        """
        Validate that `batch_size` is not set in the dataloader parameters.

        `batch_size` must be set through `batch_size` field, not
        through the dataloader parameters.

        Parameters
        ----------
        dataloader_params : dict of {str: Any}
            The dataloader parameters.

        Returns
        -------
        dict of {str: Any}
            The validated dataloader parameters.

        Raises
        ------
        ValueError
            If `batch_size` is present in the dataloader parameters.
        """
        if "batch_size" in dataloader_params:
            raise ValueError(
                "`batch_size` should not be set in the dataloader parameters. "
                "Use the `batch_size` field of `NGDataConfig` instead."
            )
        return dataloader_params

    @field_validator("train_dataloader_params")
    @classmethod
    def shuffle_train_dataloader(
        cls, train_dataloader_params: dict[str, Any]
    ) -> dict[str, Any]:
        """
        Validate that "shuffle" is included in the training dataloader params.

        A warning will be raised if `shuffle=False`.

        Parameters
        ----------
        train_dataloader_params : dict of {str: Any}
            The training dataloader parameters.

        Returns
        -------
        dict of {str: Any}
            The validated training dataloader parameters.

        Raises
        ------
        ValueError
            If "shuffle" is not included in the training dataloader params.
        """
        if "shuffle" not in train_dataloader_params:
            raise ValueError(
                "`train_dataloader_params` must include the `shuffle` parameter."
            )
        elif ("shuffle" in train_dataloader_params) and (
            not train_dataloader_params["shuffle"]
        ):
            warn(
                "`train_dataloader_params` includes `shuffle=False`, which may lead to "
                "lower quality results.",
                stacklevel=1,
            )
        return train_dataloader_params

    @model_validator(mode="after")
    def validate_dimensions(self: Self) -> Self:
        """
        Validate 2D/3D dimensions between axes and patch size.

        Returns
        -------
        Self
            Validated data model.

        Raises
        ------
        ValueError
            If the patch size dimension is not compatible with the axes.
        """
        # "whole" patching does not have dimensions to validate
        if not hasattr(self.patching, "patch_size"):
            return self

        if self.data_type == "czi":
            # Z and T are both depth axes for CZI data
            expected_dims = 3 if ("Z" in self.axes or "T" in self.axes) else 2
            additional_message = " (`Z` and `T` are depth axes for CZI data)"
        else:
            expected_dims = 3 if "Z" in self.axes else 2
            additional_message = ""

        # infer dimension from requested patch size
        actual_dims = len(self.patching.patch_size)
        if actual_dims != expected_dims:
            raise ValueError(
                f"`patch_size` in `patching` must have {expected_dims} dimensions, "
                f"got {self.patching.patch_size} with axes {self.axes}"
                f"{additional_message}."
            )

        return self

    @model_validator(mode="after")
    def propagate_seed_to_filters(self: Self) -> Self:
        """
        Propagate the main seed to patch and coordinate filters that support seeds.

        This ensures that all filters use the same seed for reproducibility,
        unless they already have a seed explicitly set.

        Returns
        -------
        Self
            Data model with propagated seeds.
        """
        if self.seed is not None:
            if self.patch_filter is not None:
                if (
                    hasattr(self.patch_filter, "seed")
                    and self.patch_filter.seed is None
                ):
                    self.patch_filter.seed = self.seed

            if self.coord_filter is not None:
                if (
                    hasattr(self.coord_filter, "seed")
                    and self.coord_filter.seed is None
                ):
                    self.coord_filter.seed = self.seed

        return self

    @model_validator(mode="after")
    def propagate_seed_to_augmentations(self: Self) -> Self:
        """
        Propagate the main seed to all augmentations that support seeds.

        This ensures that all augmentations use the same seed for
         reproducibility, unless they already have a seed explicitly set.

        Returns
        -------
        Self
            Data model with propagated seeds.
        """
        if self.seed is not None:
            for transform in self.augmentations:
                if hasattr(transform, "seed") and transform.seed is None:
                    transform.seed = self.seed
        return self

    @model_validator(mode="after")
    def propagate_seed_to_patching(self: Self) -> Self:
        """
        Propagate the main seed to the patching strategy if it supports seeds.

        This ensures that the patching strategy uses the same seed for reproducibility,
        unless it already has a seed explicitly set.

        Returns
        -------
        Self
            Data model with propagated seed.
        """
        if self.seed is not None:
            if hasattr(self.patching, "seed") and self.patching.seed is None:
                self.patching.seed = self.seed
        return self

    @field_validator("train_dataloader_params", "val_dataloader_params", mode="before")
    @classmethod
    def set_default_pin_memory(
        cls, dataloader_params: dict[str, Any]
    ) -> dict[str, Any]:
        """
        Set default pin_memory for dataloader parameters if not provided.

        - If 'pin_memory' is not set, it defaults to True if CUDA is available.

        Parameters
        ----------
        dataloader_params : dict of {str: Any}
            The dataloader parameters.

        Returns
        -------
        dict of {str: Any}
            The dataloader parameters with pin_memory default applied.
        """
        if "pin_memory" not in dataloader_params:
            import torch

            dataloader_params["pin_memory"] = torch.cuda.is_available()
        return dataloader_params

    @field_validator("train_dataloader_params", mode="before")
    @classmethod
    def set_default_train_workers(
        cls, dataloader_params: dict[str, Any]
    ) -> dict[str, Any]:
        """
        Set default num_workers for training dataloader if not provided.

        - If 'num_workers' is not set, it defaults to the number of available CPU cores.

        Parameters
        ----------
        dataloader_params : dict of {str: Any}
            The training dataloader parameters.

        Returns
        -------
        dict of {str: Any}
            The dataloader parameters with num_workers default applied.
        """
        if "num_workers" not in dataloader_params:
            # Use 0 workers during tests, otherwise use all available CPU cores
            if "pytest" in sys.modules:
                dataloader_params["num_workers"] = 0
            else:
                dataloader_params["num_workers"] = os.cpu_count()

        return dataloader_params

    @model_validator(mode="after")
    def set_val_workers_to_match_train(self: Self) -> Self:
        """
        Set validation dataloader num_workers to match training dataloader.

        If num_workers is not specified in val_dataloader_params, it will be set to the
        same value as train_dataloader_params["num_workers"].

        Returns
        -------
        Self
            Validated data model with synchronized num_workers.
        """
        if "num_workers" not in self.val_dataloader_params:
            self.val_dataloader_params["num_workers"] = self.train_dataloader_params[
                "num_workers"
            ]
        return self

    def __str__(self) -> str:
        """
        Pretty string reprensenting the configuration.

        Returns
        -------
        str
            Pretty string.
        """
        return pformat(self.model_dump())

    def is_3D(self) -> bool:
        """
        Check if the data is 3D based on the axes.

        Either "Z" is in the axes and patching `patch_size` has 3 dimensions, or for CZI
        data, "Z" is in the axes or "T" is in the axes and patching `patch_size` has
        3 dimensions.

        This method is used during NGConfiguration validation to cross checks dimensions
        with the algorithm configuration.

        Returns
        -------
        bool
            True if the data is 3D, False otherwise.
        """
        if self.data_type == "czi":
            return "Z" in self.axes or "T" in self.axes
        else:
            return "Z" in self.axes

    # TODO: if switching from a state in which in_memory=True to an incompatible state
    # an error will be raised. Should that automatically be set to False instead?
    # TODO this method could be private and we could have public `to_validation_config`
    #   and `to_prediction_config` methods with appropriate parameters
    def convert_mode(
        self,
        new_mode: Literal["validating", "predicting"],
        new_patch_size: Sequence[int] | None = None,
        overlap_size: Sequence[int] | None = None,
        new_batch_size: int | None = None,
        new_data_type: Literal["array", "tiff", "zarr", "czi", "custom"] | None = None,
        new_axes: str | None = None,
        new_channels: Sequence[int] | Literal["all"] | None = None,
        new_in_memory: bool | None = None,
        new_dataloader_params: dict[str, Any] | None = None,
    ) -> NGDataConfig:
        """
        Convert a training dataset configuration to a different mode.

        This method is intended to facilitate creating validation or prediction
        configurations from a training configuration.

        To perform tile prediction when switching to `predicting` mode, please provide
        both `new_patch_size` and `overlap_size`. Switching mode to `predicting` without
        specifying `new_patch_size` and `overlap_size` will apply the default patching
        strategy, namely `whole` image strategy. `new_patch_size` and `overlap_size` are
        only used when switching to `predicting`.

        `channels=None` will retain the same channels as in the current configuration.
        To select all channels, please specify all channels explicitly or pass
        `channels='all'`.

        New dataloader parameters will be placed in the appropriate dataloader params
        field depending on the new mode.

        To create a new training configuration, please use
        `careamics.config.create_ng_data_configuration`.

        This method compares the new parameters with the current ones and raises
        errors if incompatible changes are requested, such as switching between 2D and
        3D axes, or changing the number of channels. Incompatibility across parameters
        may be delegated to Pydantic validation.

        Parameters
        ----------
        new_mode : Literal["validating", "predicting"]
            The new dataset mode, one of `validating` or `predicting`.
        new_patch_size : Sequence of int, default=None
            New patch size. If None for `predicting`, uses default whole image strategy.
        overlap_size : Sequence of int, default=None
            New overlap size. Necessary when switching to `predicting` with tiled
            patching.
        new_batch_size : int, default=None
            New batch size.
        new_data_type : Literal['array', 'tiff', 'zarr', 'czi', 'custom'], default=None
            New data type.
        new_axes : str, default=None
            New axes.
        new_channels : Sequence of int or "all", default=None
            New channels.
        new_in_memory : bool, default=None
            New in_memory value.
        new_dataloader_params : dict of {str: Any}, default=None
            New dataloader parameters. These will be placed in the
            appropriate dataloader params field depending on the new mode.

        Returns
        -------
        NGDataConfig
            New NGDataConfig with the updated mode and parameters.

        Raises
        ------
        ValueError
            If conversion to training mode is requested, or if incompatible changes
            are requested.
        """
        if self.mode != Mode.TRAINING:
            raise ValueError(
                f"Conversion from mode '{self.mode}' to '{new_mode}' is not supported. "
                f"Only conversion from 'training' mode is supported."
            )
        if new_mode == Mode.TRAINING:
            raise ValueError(
                "Conversion to 'training' mode is not supported. Create a new "
                "NGDataConfig instead, for instance using "
                "`create_ng_data_configuration`."
            )

        # sanity checks
        # switching spatial axes
        if not _are_spatial_dims_maintained(
            self.data_type,
            self.axes,
            new_data_type or self.data_type,
            new_axes or self.axes,
        ):  # switching 2D/3D
            additional_msg = ""
            if self.data_type == "czi" or new_data_type == "czi":
                additional_msg = " Note that for CZI data, Z and T are both depth axes."

            raise ValueError(
                "Conversion between different spatial dimensions is not allowed. Got "
                f"new axes {new_axes} with new data type {new_data_type}, and current "
                f"axes {self.axes} with current data type {self.data_type}."
                f"{additional_msg}"
            )

        # normalize new_channels parameter to lift ambiguity around `None`
        #   - If None, keep previous parameter
        #   - If "all", select all channels (None value internally)
        if new_channels is None:
            new_channels = self.channels
        elif new_channels == "all":
            new_channels = None  # all channels

        # switching channels
        _validate_channel_conversion(
            self.axes,
            self.channels,
            new_axes or self.axes,  # if new_axes is None, we keep the same axes
            new_channels,  # new_channel has already been updated to the correct value
        )

        # apply default values
        patching_strategy: PatchingConfig
        if new_mode == Mode.PREDICTING:
            if new_patch_size is None:
                patching_strategy = WholePatchingConfig()
            else:
                if overlap_size is None:
                    raise ValueError(
                        "`overlap_size` parameter must be specified when switching to "
                        "'predicting' mode with a `new_patch_size`."
                    )
                patching_strategy = TiledPatchingConfig(
                    patch_size=list(new_patch_size), overlaps=list(overlap_size)
                )
        else:  # validating
            # to satisfy mypy, since self.mode=="training", patching has patch_size
            assert not isinstance(self.patching, WholePatchingConfig)

            patching_strategy = FixedRandomPatchingConfig(
                patch_size=(
                    list(new_patch_size)
                    if new_patch_size is not None
                    else self.patching.patch_size
                ),
            )

        # create new config
        model_dict = self.model_dump()
        model_dict.update(
            {
                "mode": new_mode,
                "patching": patching_strategy,
                "batch_size": new_batch_size or self.batch_size,
                "data_type": new_data_type or self.data_type,
                "axes": new_axes or self.axes,
                "channels": new_channels,
                "in_memory": (
                    new_in_memory if new_in_memory is not None else self.in_memory
                ),
                "val_dataloader_params": (
                    new_dataloader_params
                    if new_mode == Mode.VALIDATING and new_dataloader_params is not None
                    else self.val_dataloader_params
                ),
                "pred_dataloader_params": (
                    new_dataloader_params
                    if new_mode == Mode.PREDICTING and new_dataloader_params is not None
                    else self.pred_dataloader_params
                ),
            }
        )

        # remove patch and coord filters when switching to validation or prediction
        del model_dict["patch_filter"]
        del model_dict["coord_filter"]

        return NGDataConfig(**model_dict)