configuration_model

class Configuration(BaseModel):
    """
    CAREamics configuration.

    The configuration defines all parameters used to build and train a CAREamics model.
    These parameters are validated to ensure that they are compatible with each other.

    It contains three sub-configurations:

    - AlgorithmModel: configuration for the algorithm training, which includes the
        architecture, loss function, optimizer, and other hyperparameters.
    - DataModel: configuration for the dataloader, which includes the type of data,
        transformations, mean/std and other parameters.
    - TrainingModel: configuration for the training, which includes the number of
        epochs or the callbacks.

    Attributes
    ----------
    experiment_name : str
        Name of the experiment, used when saving logs and checkpoints.
    algorithm : AlgorithmModel
        Algorithm configuration.
    data : DataModel
        Data configuration.
    training : TrainingModel
        Training configuration.

    Methods
    -------
    set_3D(is_3D: bool, axes: str, patch_size: List[int]) -> None
        Switch configuration between 2D and 3D.
    set_N2V2(use_n2v2: bool) -> None
        Switch N2V algorithm between N2V and N2V2.
    set_structN2V(
        mask_axis: Literal["horizontal", "vertical", "none"], mask_span: int) -> None
        Set StructN2V parameters.
    model_dump(
        exclude_defaults: bool = False, exclude_none: bool = True, **kwargs: Dict
        ) -> Dict
        Export configuration to a dictionary.

    Raises
    ------
    ValueError
        Configuration parameter type validation errors.
    ValueError
        If the experiment name contains invalid characters or is empty.
    ValueError
        If the algorithm is 3D but there is not "Z" in the data axes, or 2D algorithm
        with "Z" in data axes.
    ValueError
        Algorithm, data or training validation errors.

    Notes
    -----
    We provide convenience methods to create standards configurations, for instance
    for N2V, in the `careamics.config.configuration_factory` module.
    >>> from careamics.config.configuration_factory import create_n2v_configuration
    >>> config = create_n2v_configuration(
    ...     experiment_name="n2v_experiment",
    ...     data_type="array",
    ...     axes="YX",
    ...     patch_size=[64, 64],
    ...     batch_size=32,
    ...     num_epochs=100
    ... )

    The configuration can be exported to a dictionary using the model_dump method:
    >>> config_dict = config.model_dump()

    Configurations can also be exported or imported from yaml files:
    >>> from careamics.config import save_configuration, load_configuration
    >>> path_to_config = save_configuration(config, my_path / "config.yml")
    >>> other_config = load_configuration(path_to_config)

    Examples
    --------
    Minimum example:
    >>> from careamics.config import Configuration
    >>> config_dict = {
    ...         "experiment_name": "N2V_experiment",
    ...         "algorithm_config": {
    ...             "algorithm": "n2v",
    ...             "loss": "n2v",
    ...             "model": {
    ...                 "architecture": "UNet",
    ...             },
    ...         },
    ...         "training_config": {
    ...             "num_epochs": 200,
    ...         },
    ...         "data_config": {
    ...             "data_type": "tiff",
    ...             "patch_size": [64, 64],
    ...             "axes": "SYX",
    ...         },
    ...     }
    >>> config = Configuration(**config_dict)
    """

    model_config = ConfigDict(
        validate_assignment=True,
        set_arbitrary_types_allowed=True,
    )

    # version
    version: Literal["0.1.0"] = "0.1.0"
    """CAREamics configuration version."""

    # required parameters
    experiment_name: str
    """Name of the experiment, used to name logs and checkpoints."""

    # Sub-configurations
    algorithm_config: Union[FCNAlgorithmConfig, VAEAlgorithmConfig] = Field(
        discriminator="algorithm"
    )
    """Algorithm configuration, holding all parameters required to configure the
    model."""

    data_config: DataConfig
    """Data configuration, holding all parameters required to configure the training
    data loader."""

    training_config: TrainingConfig
    """Training configuration, holding all parameters required to configure the
    training process."""

    @field_validator("experiment_name")
    @classmethod
    def no_symbol(cls, name: str) -> str:
        """
        Validate experiment name.

        A valid experiment name is a non-empty string with only contains letters,
        numbers, underscores, dashes and spaces.

        Parameters
        ----------
        name : str
            Name to validate.

        Returns
        -------
        str
            Validated name.

        Raises
        ------
        ValueError
            If the name is empty or contains invalid characters.
        """
        if len(name) == 0 or name.isspace():
            raise ValueError("Experiment name is empty.")

        # Validate using a regex that it contains only letters, numbers, underscores,
        # dashes and spaces
        if not re.match(r"^[a-zA-Z0-9_\- ]*$", name):
            raise ValueError(
                f"Experiment name contains invalid characters (got {name}). "
                f"Only letters, numbers, underscores, dashes and spaces are allowed."
            )

        return name

    @model_validator(mode="after")
    def validate_3D(self: Self) -> Self:
        """
        Change algorithm dimensions to match data.axes.

        Only for non-custom algorithms.

        Returns
        -------
        Self
            Validated configuration.
        """
        if self.algorithm_config.algorithm != SupportedAlgorithm.CUSTOM:
            if "Z" in self.data_config.axes and not self.algorithm_config.model.is_3D():
                # change algorithm to 3D
                self.algorithm_config.model.set_3D(True)
            elif (
                "Z" not in self.data_config.axes and self.algorithm_config.model.is_3D()
            ):
                # change algorithm to 2D
                self.algorithm_config.model.set_3D(False)

        return self

    @model_validator(mode="after")
    def validate_algorithm_and_data(self: Self) -> Self:
        """
        Validate algorithm and data compatibility.

        In particular, the validation does the following:

        - If N2V is used, it enforces the presence of N2V_Maniuplate in the transforms
        - If N2V2 is used, it enforces the correct manipulation strategy

        Returns
        -------
        Self
            Validated configuration.
        """
        if self.algorithm_config.algorithm == SupportedAlgorithm.N2V:
            # missing N2V_MANIPULATE
            if not self.data_config.has_n2v_manipulate():
                self.data_config.transforms.append(
                    N2VManipulateModel(
                        name=SupportedTransform.N2V_MANIPULATE.value,
                    )
                )

            median = SupportedPixelManipulation.MEDIAN.value
            uniform = SupportedPixelManipulation.UNIFORM.value
            strategy = median if self.algorithm_config.model.n2v2 else uniform
            self.data_config.set_N2V2_strategy(strategy)
        else:
            # remove N2V manipulate if present
            if self.data_config.has_n2v_manipulate():
                self.data_config.remove_n2v_manipulate()

        return self

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

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

    def set_3D(self, is_3D: bool, axes: str, patch_size: list[int]) -> None:
        """
        Set 3D flag and axes.

        Parameters
        ----------
        is_3D : bool
            Whether the algorithm is 3D or not.
        axes : str
            Axes of the data.
        patch_size : list[int]
            Patch size.
        """
        # set the flag and axes (this will not trigger validation at the config level)
        self.algorithm_config.model.set_3D(is_3D)
        self.data_config.set_3D(axes, patch_size)

        # cheap hack: trigger validation
        self.algorithm_config = self.algorithm_config

    def set_N2V2(self, use_n2v2: bool) -> None:
        """
        Switch N2V algorithm between N2V and N2V2.

        Parameters
        ----------
        use_n2v2 : bool
            Whether to use N2V2 or not.

        Raises
        ------
        ValueError
            If the algorithm is not N2V.
        """
        if self.algorithm_config.algorithm == SupportedAlgorithm.N2V:
            self.algorithm_config.model.n2v2 = use_n2v2
            strategy = (
                SupportedPixelManipulation.MEDIAN.value
                if use_n2v2
                else SupportedPixelManipulation.UNIFORM.value
            )
            self.data_config.set_N2V2_strategy(strategy)
        else:
            raise ValueError("N2V2 can only be set for N2V algorithm.")

    def set_structN2V(
        self, mask_axis: Literal["horizontal", "vertical", "none"], mask_span: int
    ) -> None:
        """
        Set StructN2V parameters.

        Parameters
        ----------
        mask_axis : Literal["horizontal", "vertical", "none"]
            Axis of the structural mask.
        mask_span : int
            Span of the structural mask.
        """
        self.data_config.set_structN2V_mask(mask_axis, mask_span)

    def get_algorithm_flavour(self) -> str:
        """
        Get the algorithm name.

        Returns
        -------
        str
            Algorithm name.
        """
        if self.algorithm_config.algorithm == SupportedAlgorithm.N2V:
            use_n2v2 = self.algorithm_config.model.n2v2
            use_structN2V = self.data_config.transforms[-1].struct_mask_axis != "none"

            # return the n2v flavour
            if use_n2v2 and use_structN2V:
                return STRUCT_N2V2
            elif use_n2v2:
                return N2V2
            elif use_structN2V:
                return STRUCT_N2V
            else:
                return N2V
        elif self.algorithm_config.algorithm == SupportedAlgorithm.N2N:
            return N2N
        elif self.algorithm_config.algorithm == SupportedAlgorithm.CARE:
            return CARE
        else:
            return CUSTOM

    def get_algorithm_description(self) -> str:
        """
        Return a description of the algorithm.

        This method is used to generate the README of the BioImage Model Zoo export.

        Returns
        -------
        str
            Description of the algorithm.
        """
        algorithm_flavour = self.get_algorithm_flavour()

        if algorithm_flavour == CUSTOM:
            return f"Custom algorithm, named {self.algorithm_config.model.name}"
        else:  # currently only N2V flavours
            if algorithm_flavour == N2V:
                return N2VDescription().description
            elif algorithm_flavour == N2V2:
                return N2V2Description().description
            elif algorithm_flavour == STRUCT_N2V:
                return StructN2VDescription().description
            elif algorithm_flavour == STRUCT_N2V2:
                return StructN2V2Description().description
            elif algorithm_flavour == N2N:
                return N2NDescription().description
            elif algorithm_flavour == CARE:
                return CAREDescription().description

        return ""

    def get_algorithm_citations(self) -> list[CiteEntry]:
        """
        Return a list of citation entries of the current algorithm.

        This is used to generate the model description for the BioImage Model Zoo.

        Returns
        -------
        List[CiteEntry]
            List of citation entries.
        """
        if self.algorithm_config.algorithm == SupportedAlgorithm.N2V:
            use_n2v2 = self.algorithm_config.model.n2v2
            use_structN2V = self.data_config.transforms[-1].struct_mask_axis != "none"

            # return the (struct)N2V(2) references
            if use_n2v2 and use_structN2V:
                return [N2VRef, N2V2Ref, StructN2VRef]
            elif use_n2v2:
                return [N2VRef, N2V2Ref]
            elif use_structN2V:
                return [N2VRef, StructN2VRef]
            else:
                return [N2VRef]
        elif self.algorithm_config.algorithm == SupportedAlgorithm.N2N:
            return [N2NRef]
        elif self.algorithm_config.algorithm == SupportedAlgorithm.CARE:
            return [CARERef]

        raise ValueError("Citation not available for custom algorithm.")

    def get_algorithm_references(self) -> str:
        """
        Get the algorithm references.

        This is used to generate the README of the BioImage Model Zoo export.

        Returns
        -------
        str
            Algorithm references.
        """
        if self.algorithm_config.algorithm == SupportedAlgorithm.N2V:
            use_n2v2 = self.algorithm_config.model.n2v2
            use_structN2V = self.data_config.transforms[-1].struct_mask_axis != "none"

            references = [
                N2VRef.text + " doi: " + N2VRef.doi,
                N2V2Ref.text + " doi: " + N2V2Ref.doi,
                StructN2VRef.text + " doi: " + StructN2VRef.doi,
            ]

            # return the (struct)N2V(2) references
            if use_n2v2 and use_structN2V:
                return "".join(references)
            elif use_n2v2:
                references.pop(-1)
                return "".join(references)
            elif use_structN2V:
                references.pop(-2)
                return "".join(references)
            else:
                return references[0]

        return ""

    def get_algorithm_keywords(self) -> list[str]:
        """
        Get algorithm keywords.

        Returns
        -------
        list[str]
            List of keywords.
        """
        if self.algorithm_config.algorithm == SupportedAlgorithm.N2V:
            use_n2v2 = self.algorithm_config.model.n2v2
            use_structN2V = self.data_config.transforms[-1].struct_mask_axis != "none"

            keywords = [
                "denoising",
                "restoration",
                "UNet",
                "3D" if "Z" in self.data_config.axes else "2D",
                "CAREamics",
                "pytorch",
                N2V,
            ]

            if use_n2v2:
                keywords.append(N2V2)
            if use_structN2V:
                keywords.append(STRUCT_N2V)
        else:
            keywords = ["CAREamics"]

        return keywords

    def model_dump(
        self,
        exclude_defaults: bool = False,
        exclude_none: bool = True,
        **kwargs: dict,
    ) -> dict:
        """
        Override model_dump method in order to set default values.

        Parameters
        ----------
        exclude_defaults : bool, optional
            Whether to exclude fields with default values or not, by default
            True.
        exclude_none : bool, optional
            Whether to exclude fields with None values or not, by default True.
        **kwargs : dict
            Keyword arguments.

        Returns
        -------
        dict
            Dictionary containing the model parameters.
        """
        dictionary = super().model_dump(
            exclude_none=exclude_none, exclude_defaults=exclude_defaults, **kwargs
        )

        return dictionary