configuration_factories

def algorithm_factory(
    algorithm: dict[str, Any],
) -> Union[N2VAlgorithm, N2NAlgorithm, CAREAlgorithm]:
    """
    Create an algorithm model for training CAREamics.

    Parameters
    ----------
    algorithm : dict
        Algorithm dictionary.

    Returns
    -------
    N2VAlgorithm or N2NAlgorithm or CAREAlgorithm
        Algorithm model for training CAREamics.
    """
    adapter: TypeAdapter = TypeAdapter(
        Annotated[
            Union[N2VAlgorithm, N2NAlgorithm, CAREAlgorithm],
            Field(discriminator="algorithm"),
        ]
    )
    return adapter.validate_python(algorithm)

def create_care_configuration(
    experiment_name: str,
    data_type: Literal["array", "tiff", "czi", "custom"],
    axes: str,
    patch_size: Sequence[int],
    batch_size: int,
    num_epochs: int = 100,
    num_steps: int | None = None,
    augmentations: list[Union[XYFlipModel, XYRandomRotate90Model]] | None = None,
    independent_channels: bool = True,
    loss: Literal["mae", "mse"] = "mae",
    n_channels_in: int | None = None,
    n_channels_out: int | None = None,
    logger: Literal["wandb", "tensorboard", "none"] = "none",
    trainer_params: dict | None = None,
    model_params: dict | None = None,
    optimizer: Literal["Adam", "Adamax", "SGD"] = "Adam",
    optimizer_params: dict[str, Any] | None = None,
    lr_scheduler: Literal["ReduceLROnPlateau", "StepLR"] = "ReduceLROnPlateau",
    lr_scheduler_params: dict[str, Any] | None = None,
    train_dataloader_params: dict[str, Any] | None = None,
    val_dataloader_params: dict[str, Any] | None = None,
    checkpoint_params: dict[str, Any] | None = None,
) -> Configuration:
    """
    Create a configuration for training CARE.

    If "Z" is present in `axes`, then `path_size` must be a list of length 3, otherwise
    2.

    If "C" is present in `axes`, then you need to set `n_channels_in` to the number of
    channels. Likewise, if you set the number of channels, then "C" must be present in
    `axes`.

    To set the number of output channels, use the `n_channels_out` parameter. If it is
    not specified, it will be assumed to be equal to `n_channels_in`.

    By default, all channels are trained together. To train all channels independently,
    set `independent_channels` to True.

    By setting `augmentations` to `None`, the default transformations (flip in X and Y,
    rotations by 90 degrees in the XY plane) are applied. Rather than the default
    transforms, a list of transforms can be passed to the `augmentations` parameter. To
    disable the transforms, simply pass an empty list.

    Parameters
    ----------
    experiment_name : str
        Name of the experiment.
    data_type : Literal["array", "tiff", "czi", "custom"]
        Type of the data.
    axes : str
        Axes of the data (e.g. SYX).
    patch_size : List[int]
        Size of the patches along the spatial dimensions (e.g. [64, 64]).
    batch_size : int
        Batch size.
    num_epochs : int, default=100
        Number of epochs to train for. If provided, this will be added to
        trainer_params.
    num_steps : int, optional
        Number of batches in 1 epoch. If provided, this will be added to trainer_params.
        Translates to `limit_train_batches` in PyTorch Lightning Trainer. See relevant
        documentation for more details.
    augmentations : list of transforms, default=None
        List of transforms to apply, either both or one of XYFlipModel and
        XYRandomRotate90Model. By default, it applies both XYFlip (on X and Y)
        and XYRandomRotate90 (in XY) to the images.
    independent_channels : bool, optional
        Whether to train all channels independently, by default False.
    loss : Literal["mae", "mse"], default="mae"
        Loss function to use.
    n_channels_in : int or None, default=None
        Number of channels in.
    n_channels_out : int or None, default=None
        Number of channels out.
    logger : Literal["wandb", "tensorboard", "none"], default="none"
        Logger to use.
    trainer_params : dict, optional
        Parameters for the trainer class, see PyTorch Lightning documentation.
    model_params : dict, default=None
        UNetModel parameters.
    optimizer : Literal["Adam", "Adamax", "SGD"], default="Adam"
        Optimizer to use.
    optimizer_params : dict, default=None
        Parameters for the optimizer, see PyTorch documentation for more details.
    lr_scheduler : Literal["ReduceLROnPlateau", "StepLR"], default="ReduceLROnPlateau"
        Learning rate scheduler to use.
    lr_scheduler_params : dict, default=None
        Parameters for the learning rate scheduler, see PyTorch documentation for more
        details.
    train_dataloader_params : dict, optional
        Parameters for the training dataloader, see the PyTorch docs for `DataLoader`.
        If left as `None`, the dict `{"shuffle": True}` will be used, this is set in
        the `GeneralDataConfig`.
    val_dataloader_params : dict, optional
        Parameters for the validation dataloader, see PyTorch the docs for `DataLoader`.
        If left as `None`, the empty dict `{}` will be used, this is set in the
        `GeneralDataConfig`.
    checkpoint_params : dict, default=None
        Parameters for the checkpoint callback, see PyTorch Lightning documentation
        (`ModelCheckpoint`) for the list of available parameters.

    Returns
    -------
    Configuration
        Configuration for training CARE.

    Examples
    --------
    Minimum example:
    >>> config = create_care_configuration(
    ...     experiment_name="care_experiment",
    ...     data_type="array",
    ...     axes="YX",
    ...     patch_size=[64, 64],
    ...     batch_size=32,
    ...     num_epochs=100
    ... )

    You can also limit the number of batches per epoch:
    >>> config = create_care_configuration(
    ...     experiment_name="care_experiment",
    ...     data_type="array",
    ...     axes="YX",
    ...     patch_size=[64, 64],
    ...     batch_size=32,
    ...     num_steps=100  # limit to 100 batches per epoch
    ... )

    To disable transforms, simply set `augmentations` to an empty list:
    >>> config = create_care_configuration(
    ...     experiment_name="care_experiment",
    ...     data_type="array",
    ...     axes="YX",
    ...     patch_size=[64, 64],
    ...     batch_size=32,
    ...     num_epochs=100,
    ...     augmentations=[]
    ... )

    A list of transforms can be passed to the `augmentations` parameter to replace the
    default augmentations:
    >>> from careamics.config.transformations import XYFlipModel
    >>> config = create_care_configuration(
    ...     experiment_name="care_experiment",
    ...     data_type="array",
    ...     axes="YX",
    ...     patch_size=[64, 64],
    ...     batch_size=32,
    ...     num_epochs=100,
    ...     augmentations=[
    ...         # No rotation and only Y flipping
    ...         XYFlipModel(flip_x = False, flip_y = True)
    ...     ]
    ... )

    If you are training multiple channels they will be trained independently by default,
    you simply need to specify the number of channels input (and optionally, the number
    of channels output):
    >>> config = create_care_configuration(
    ...     experiment_name="care_experiment",
    ...     data_type="array",
    ...     axes="YXC", # channels must be in the axes
    ...     patch_size=[64, 64],
    ...     batch_size=32,
    ...     num_epochs=100,
    ...     n_channels_in=3, # number of input channels
    ...     n_channels_out=1 # if applicable
    ... )

    If instead you want to train multiple channels together, you need to turn off the
    `independent_channels` parameter:
    >>> config = create_care_configuration(
    ...     experiment_name="care_experiment",
    ...     data_type="array",
    ...     axes="YXC", # channels must be in the axes
    ...     patch_size=[64, 64],
    ...     batch_size=32,
    ...     num_epochs=100,
    ...     independent_channels=False,
    ...     n_channels_in=3,
    ...     n_channels_out=1 # if applicable
    ... )

    If you would like to train on CZI files, use `"czi"` as `data_type` and `"SCYX"` as
    `axes` for 2-D or `"SCZYX"` for 3-D denoising. Note that `"SCYX"` can also be used
    for 3-D data but spatial context along the Z dimension will then not be taken into
    account.
    >>> config_2d = create_care_configuration(
    ...     experiment_name="care_experiment",
    ...     data_type="czi",
    ...     axes="SCYX",
    ...     patch_size=[64, 64],
    ...     batch_size=32,
    ...     num_epochs=100,
    ...     n_channels_in=1,
    ... )
    >>> config_3d = create_care_configuration(
    ...     experiment_name="care_experiment",
    ...     data_type="czi",
    ...     axes="SCZYX",
    ...     patch_size=[16, 64, 64],
    ...     batch_size=16,
    ...     num_epochs=100,
    ...     n_channels_in=1,
    ... )
    """
    return Configuration(
        **_create_supervised_config_dict(
            algorithm="care",
            experiment_name=experiment_name,
            data_type=data_type,
            axes=axes,
            patch_size=patch_size,
            batch_size=batch_size,
            augmentations=augmentations,
            independent_channels=independent_channels,
            loss=loss,
            n_channels_in=n_channels_in,
            n_channels_out=n_channels_out,
            logger=logger,
            trainer_params=trainer_params,
            model_params=model_params,
            optimizer=optimizer,
            optimizer_params=optimizer_params,
            lr_scheduler=lr_scheduler,
            lr_scheduler_params=lr_scheduler_params,
            train_dataloader_params=train_dataloader_params,
            val_dataloader_params=val_dataloader_params,
            checkpoint_params=checkpoint_params,
            num_epochs=num_epochs,
            num_steps=num_steps,
        )
    )

def create_hdn_configuration(
    experiment_name: str,
    data_type: Literal["array", "tiff", "custom"],
    axes: str,
    patch_size: Sequence[int],
    batch_size: int,
    num_epochs: int = 100,
    num_steps: int | None = None,
    encoder_conv_strides: tuple[int, ...] = (2, 2),
    decoder_conv_strides: tuple[int, ...] = (2, 2),
    multiscale_count: int = 1,
    z_dims: tuple[int, ...] = (128, 128),
    output_channels: int = 1,
    encoder_n_filters: int = 32,
    decoder_n_filters: int = 32,
    encoder_dropout: float = 0.0,
    decoder_dropout: float = 0.0,
    nonlinearity: Literal[
        "None", "Sigmoid", "Softmax", "Tanh", "ReLU", "LeakyReLU", "ELU"
    ] = "ReLU",
    analytical_kl: bool = False,
    predict_logvar: Literal["pixelwise"] | None = None,
    logvar_lowerbound: Union[float, None] = None,
    logger: Literal["wandb", "tensorboard", "none"] = "none",
    trainer_params: dict | None = None,
    augmentations: list[Union[XYFlipModel, XYRandomRotate90Model]] | None = None,
    train_dataloader_params: dict[str, Any] | None = None,
    val_dataloader_params: dict[str, Any] | None = None,
) -> Configuration:
    """
    Create a configuration for training HDN.

    If "Z" is present in `axes`, then `path_size` must be a list of length 3, otherwise
    2.

    If "C" is present in `axes`, then you need to set `n_channels_in` to the number of
    channels. Likewise, if you set the number of channels, then "C" must be present in
    `axes`.

    To set the number of output channels, use the `n_channels_out` parameter. If it is
    not specified, it will be assumed to be equal to `n_channels_in`.

    By default, all channels are trained independently. To train all channels together,
    set `independent_channels` to False.

    By setting `augmentations` to `None`, the default transformations (flip in X and Y,
    rotations by 90 degrees in the XY plane) are applied. Rather than the default
    transforms, a list of transforms can be passed to the `augmentations` parameter. To
    disable the transforms, simply pass an empty list.

    # TODO revisit the necessity of model_params

    Parameters
    ----------
    experiment_name : str
        Name of the experiment.
    data_type : Literal["array", "tiff", "custom"]
        Type of the data.
    axes : str
        Axes of the data (e.g. SYX).
    patch_size : List[int]
        Size of the patches along the spatial dimensions (e.g. [64, 64]).
    batch_size : int
        Batch size.
    num_epochs : int, default=100
        Number of epochs to train for. If provided, this will be added to
        trainer_params.
    num_steps : int, optional
        Number of batches in 1 epoch. If provided, this will be added to trainer_params.
        Translates to `limit_train_batches` in PyTorch Lightning Trainer. See relevant
        documentation for more details.
    encoder_conv_strides : tuple[int, ...], optional
        Strides for the encoder convolutional layers, by default (2, 2).
    decoder_conv_strides : tuple[int, ...], optional
        Strides for the decoder convolutional layers, by default (2, 2).
    multiscale_count : int, optional
        Number of scales in the multiscale architecture, by default 1.
    z_dims : tuple[int, ...], optional
        Dimensions of the latent space, by default (128, 128).
    output_channels : int, optional
        Number of output channels, by default 1.
    encoder_n_filters : int, optional
        Number of filters in the encoder, by default 32.
    decoder_n_filters : int, optional
        Number of filters in the decoder, by default 32.
    encoder_dropout : float, optional
        Dropout rate for the encoder, by default 0.0.
    decoder_dropout : float, optional
        Dropout rate for the decoder, by default 0.0.
    nonlinearity : Literal, optional
        Nonlinearity function to use, by default "ReLU".
    analytical_kl : bool, optional
        Whether to use analytical KL divergence, by default False.
    predict_logvar : Literal[None, "pixelwise"], optional
        Type of log variance prediction, by default None.
    logvar_lowerbound : Union[float, None], optional
        Lower bound for the log variance, by default None.
    logger : Literal["wandb", "tensorboard", "none"], optional
        Logger to use for training, by default "none".
    trainer_params : dict, optional
        Parameters for the trainer class, see PyTorch Lightning documentation.
    augmentations : Optional[list[Union[XYFlipModel, XYRandomRotate90Model]]], optional
        List of augmentations to apply, by default None.
    train_dataloader_params : Optional[dict[str, Any]], optional
        Parameters for the training dataloader, by default None.
    val_dataloader_params : Optional[dict[str, Any]], optional
        Parameters for the validation dataloader, by default None.

    Returns
    -------
    Configuration
        The configuration object for training HDN.

    Examples
    --------
    Minimum example:
    >>> config = create_hdn_configuration(
    ...     experiment_name="hdn_experiment",
    ...     data_type="array",
    ...     axes="YX",
    ...     patch_size=[64, 64],
    ...     batch_size=32,
    ...     num_epochs=100
    ... )

    You can also limit the number of batches per epoch:
    >>> config = create_hdn_configuration(
    ...     experiment_name="hdn_experiment",
    ...     data_type="array",
    ...     axes="YX",
    ...     patch_size=[64, 64],
    ...     batch_size=32,
    ...     num_steps=100  # limit to 100 batches per epoch
    ... )
    """
    transform_list = _list_spatial_augmentations(augmentations)

    loss_config = LVAELossConfig(
        loss_type="hdn", denoisplit_weight=1, musplit_weight=0
    )  # TODO what are the correct defaults for HDN?

    gaussian_likelihood = GaussianLikelihoodConfig(
        predict_logvar=predict_logvar, logvar_lowerbound=logvar_lowerbound
    )

    # algorithm & model
    algorithm_params = _create_vae_based_algorithm(
        algorithm="hdn",
        loss=loss_config,
        input_shape=patch_size,
        encoder_conv_strides=encoder_conv_strides,
        decoder_conv_strides=decoder_conv_strides,
        multiscale_count=multiscale_count,
        z_dims=z_dims,
        output_channels=output_channels,
        encoder_n_filters=encoder_n_filters,
        decoder_n_filters=decoder_n_filters,
        encoder_dropout=encoder_dropout,
        decoder_dropout=decoder_dropout,
        nonlinearity=nonlinearity,
        predict_logvar=predict_logvar,
        analytical_kl=analytical_kl,
        gaussian_likelihood=gaussian_likelihood,
        nm_likelihood=None,
    )

    # data
    data_params = _create_data_configuration(
        data_type=data_type,
        axes=axes,
        patch_size=patch_size,
        batch_size=batch_size,
        augmentations=transform_list,
        train_dataloader_params=train_dataloader_params,
        val_dataloader_params=val_dataloader_params,
    )

    # Handle trainer parameters with num_epochs and num_steps
    final_trainer_params = {} if trainer_params is None else trainer_params.copy()

    # Add num_epochs and num_steps if provided
    if num_epochs is not None:
        final_trainer_params["max_epochs"] = num_epochs
    if num_steps is not None:
        final_trainer_params["limit_train_batches"] = num_steps

    # training
    training_params = _create_training_configuration(
        trainer_params=final_trainer_params,
        logger=logger,
    )

    return Configuration(
        experiment_name=experiment_name,
        algorithm_config=algorithm_params,
        data_config=data_params,
        training_config=training_params,
    )

def create_microsplit_configuration(
    experiment_name: str,
    data_type: Literal["array", "tiff", "custom"],
    axes: str,
    patch_size: Sequence[int],
    batch_size: int,
    num_epochs: int = 100,
    num_steps: int | None = None,
    encoder_conv_strides: tuple[int, ...] = (2, 2),
    decoder_conv_strides: tuple[int, ...] = (2, 2),
    multiscale_count: int = 3,
    grid_size: int = 32,  # TODO most likely can be derived from patch size
    z_dims: tuple[int, ...] = (128, 128),
    output_channels: int = 1,
    encoder_n_filters: int = 32,
    decoder_n_filters: int = 32,
    encoder_dropout: float = 0.0,
    decoder_dropout: float = 0.0,
    nonlinearity: Literal[
        "None", "Sigmoid", "Softmax", "Tanh", "ReLU", "LeakyReLU", "ELU"
    ] = "ReLU",  # TODO do we need all these?
    analytical_kl: bool = False,
    predict_logvar: Literal["pixelwise"] = "pixelwise",
    logvar_lowerbound: Union[float, None] = None,
    logger: Literal["wandb", "tensorboard", "none"] = "none",
    trainer_params: dict | None = None,
    augmentations: list[Union[XYFlipModel, XYRandomRotate90Model]] | None = None,
    nm_paths: list[str] | None = None,
    data_stats: tuple[float, float] | None = None,
    train_dataloader_params: dict[str, Any] | None = None,
    val_dataloader_params: dict[str, Any] | None = None,
) -> Configuration:
    """
    Create a configuration for training MicroSplit.

    Parameters
    ----------
    experiment_name : str
        Name of the experiment.
    data_type : Literal["array", "tiff", "custom"]
        Type of the data.
    axes : str
        Axes of the data (e.g. SYX).
    patch_size : Sequence[int]
        Size of the patches along the spatial dimensions (e.g. [64, 64]).
    batch_size : int
        Batch size.
    num_epochs : int, default=100
        Number of epochs to train for. If provided, this will be added to
        trainer_params.
    num_steps : int, optional
        Number of batches in 1 epoch. If provided, this will be added to trainer_params.
        Translates to `limit_train_batches` in PyTorch Lightning Trainer. See relevant
        documentation for more details.
    encoder_conv_strides : tuple[int, ...], optional
        Strides for the encoder convolutional layers, by default (2, 2).
    decoder_conv_strides : tuple[int, ...], optional
        Strides for the decoder convolutional layers, by default (2, 2).
    multiscale_count : int, optional
        Number of multiscale levels, by default 1.
    grid_size : int, optional
        Size of the grid for the lateral context, by default 32.
    z_dims : tuple[int, ...], optional
        List of latent dimensions for each hierarchy level in the LVAE, by default
        (128, 128).
    output_channels : int, optional
        Number of output channels for the model, by default 1.
    encoder_n_filters : int, optional
        Number of filters in the encoder, by default 32.
    decoder_n_filters : int, optional
        Number of filters in the decoder, by default 32.
    encoder_dropout : float, optional
        Dropout rate for the encoder, by default 0.0.
    decoder_dropout : float, optional
        Dropout rate for the decoder, by default 0.0.
    nonlinearity : Literal, optional
        Nonlinearity to use in the model, by default "ReLU".
    analytical_kl : bool, optional
        Whether to use analytical KL divergence, by default False.
    predict_logvar : Literal["pixelwise"] | None, optional
        Type of log-variance prediction, by default None.
    logvar_lowerbound : Union[float, None], optional
        Lower bound for the log variance, by default None.
    logger : Literal["wandb", "tensorboard", "none"], optional
        Logger to use for training, by default "none".
    trainer_params : dict, optional
        Parameters for the trainer class, see PyTorch Lightning documentation.
    augmentations : list[Union[XYFlipModel, XYRandomRotate90Model]] | None, optional
        List of augmentations to apply, by default None.
    nm_paths : list[str] | None, optional
        Paths to the noise model files, by default None.
    data_stats : tuple[float, float] | None, optional
        Data statistics (mean, std), by default None.
    train_dataloader_params : dict[str, Any] | None, optional
        Parameters for the training dataloader, by default None.
    val_dataloader_params : dict[str, Any] | None, optional
        Parameters for the validation dataloader, by default None.

    Returns
    -------
    Configuration
        A configuration object for the microsplit algorithm.

    Examples
    --------
    Minimum example:
    # >>> config = create_microsplit_configuration(
    # ...     experiment_name="microsplit_experiment",
    # ...     data_type="array",
    # ...     axes="YX",
    # ...     patch_size=[64, 64],
    # ...     batch_size=32,
    # ...     num_epochs=100

    # ... )

    # You can also limit the number of batches per epoch:
    # >>> config = create_microsplit_configuration(
    # ...     experiment_name="microsplit_experiment",
    # ...     data_type="array",
    # ...     axes="YX",
    # ...     patch_size=[64, 64],
    # ...     batch_size=32,
    # ...     num_steps=100  # limit to 100 batches per epoch
    # ... )
    """
    transform_list = _list_spatial_augmentations(augmentations)

    loss_config = LVAELossConfig(
        loss_type="denoisplit_musplit", denoisplit_weight=0.9, musplit_weight=0.1
    )  # TODO losses need to be refactored! just for example. Add validator if sum to 1

    # Create likelihood configurations
    gaussian_likelihood_config, noise_model_config, nm_likelihood_config = (
        get_likelihood_config(
            loss_type="denoisplit_musplit",
            predict_logvar=predict_logvar,
            logvar_lowerbound=logvar_lowerbound,
            nm_paths=nm_paths,
            data_stats=data_stats,
        )
    )

    # Create the LVAE model
    network_model = _create_vae_configuration(
        input_shape=patch_size,
        encoder_conv_strides=encoder_conv_strides,
        decoder_conv_strides=decoder_conv_strides,
        multiscale_count=multiscale_count,
        z_dims=z_dims,
        output_channels=output_channels,
        encoder_n_filters=encoder_n_filters,
        decoder_n_filters=decoder_n_filters,
        encoder_dropout=encoder_dropout,
        decoder_dropout=decoder_dropout,
        nonlinearity=nonlinearity,
        predict_logvar=predict_logvar,
        analytical_kl=analytical_kl,
    )

    # Create the MicroSplit algorithm configuration
    algorithm_params = {
        "algorithm": "microsplit",
        "loss": loss_config,
        "model": network_model,
        "gaussian_likelihood": gaussian_likelihood_config,
        "noise_model": noise_model_config,
        "noise_model_likelihood": nm_likelihood_config,
    }

    # Convert to MicroSplitAlgorithm instance
    algorithm_config = MicroSplitAlgorithm(**algorithm_params)

    # data
    data_params = _create_microsplit_data_configuration(
        data_type=data_type,
        axes=axes,
        patch_size=patch_size,
        grid_size=grid_size,
        multiscale_count=multiscale_count,
        batch_size=batch_size,
        augmentations=transform_list,
        train_dataloader_params=train_dataloader_params,
        val_dataloader_params=val_dataloader_params,
    )

    # Handle trainer parameters with num_epochs and num_steps
    final_trainer_params = {} if trainer_params is None else trainer_params.copy()

    # Add num_epochs and num_steps if provided
    if num_epochs is not None:
        final_trainer_params["max_epochs"] = num_epochs
    if num_steps is not None:
        final_trainer_params["limit_train_batches"] = num_steps

    # training
    training_params = _create_training_configuration(
        trainer_params=final_trainer_params,
        logger=logger,
    )

    return Configuration(
        experiment_name=experiment_name,
        algorithm_config=algorithm_config,
        data_config=data_params,
        training_config=training_params,
    )

def create_n2n_configuration(
    experiment_name: str,
    data_type: Literal["array", "tiff", "czi", "custom"],
    axes: str,
    patch_size: Sequence[int],
    batch_size: int,
    num_epochs: int = 100,
    num_steps: int | None = None,
    augmentations: list[Union[XYFlipModel, XYRandomRotate90Model]] | None = None,
    independent_channels: bool = True,
    loss: Literal["mae", "mse"] = "mae",
    n_channels_in: int | None = None,
    n_channels_out: int | None = None,
    logger: Literal["wandb", "tensorboard", "none"] = "none",
    trainer_params: dict | None = None,
    model_params: dict | None = None,
    optimizer: Literal["Adam", "Adamax", "SGD"] = "Adam",
    optimizer_params: dict[str, Any] | None = None,
    lr_scheduler: Literal["ReduceLROnPlateau", "StepLR"] = "ReduceLROnPlateau",
    lr_scheduler_params: dict[str, Any] | None = None,
    train_dataloader_params: dict[str, Any] | None = None,
    val_dataloader_params: dict[str, Any] | None = None,
    checkpoint_params: dict[str, Any] | None = None,
) -> Configuration:
    """
    Create a configuration for training Noise2Noise.

    If "Z" is present in `axes`, then `path_size` must be a list of length 3, otherwise
    2.

    If "C" is present in `axes`, then you need to set `n_channels_in` to the number of
    channels. Likewise, if you set the number of channels, then "C" must be present in
    `axes`.

    To set the number of output channels, use the `n_channels_out` parameter. If it is
    not specified, it will be assumed to be equal to `n_channels_in`.

    By default, all channels are trained together. To train all channels independently,
    set `independent_channels` to True.

    By setting `augmentations` to `None`, the default transformations (flip in X and Y,
    rotations by 90 degrees in the XY plane) are applied. Rather than the default
    transforms, a list of transforms can be passed to the `augmentations` parameter. To
    disable the transforms, simply pass an empty list.

    Parameters
    ----------
    experiment_name : str
        Name of the experiment.
    data_type : Literal["array", "tiff", "czi", "custom"]
        Type of the data.
    axes : str
        Axes of the data (e.g. SYX).
    patch_size : List[int]
        Size of the patches along the spatial dimensions (e.g. [64, 64]).
    batch_size : int
        Batch size.
    num_epochs : int, default=100
        Number of epochs to train for. If provided, this will be added to
        trainer_params.
    num_steps : int, optional
        Number of batches in 1 epoch. If provided, this will be added to trainer_params.
        Translates to `limit_train_batches` in PyTorch Lightning Trainer. See relevant
        documentation for more details.
    augmentations : list of transforms, default=None
        List of transforms to apply, either both or one of XYFlipModel and
        XYRandomRotate90Model. By default, it applies both XYFlip (on X and Y)
        and XYRandomRotate90 (in XY) to the images.
    independent_channels : bool, optional
        Whether to train all channels independently, by default False.
    loss : Literal["mae", "mse"], optional
        Loss function to use, by default "mae".
    n_channels_in : int or None, default=None
        Number of channels in.
    n_channels_out : int or None, default=None
        Number of channels out.
    logger : Literal["wandb", "tensorboard", "none"], optional
        Logger to use, by default "none".
    trainer_params : dict, optional
        Parameters for the trainer class, see PyTorch Lightning documentation.
    model_params : dict, default=None
        UNetModel parameters.
    optimizer : Literal["Adam", "Adamax", "SGD"], default="Adam"
        Optimizer to use.
    optimizer_params : dict, default=None
        Parameters for the optimizer, see PyTorch documentation for more details.
    lr_scheduler : Literal["ReduceLROnPlateau", "StepLR"], default="ReduceLROnPlateau"
        Learning rate scheduler to use.
    lr_scheduler_params : dict, default=None
        Parameters for the learning rate scheduler, see PyTorch documentation for more
        details.
    train_dataloader_params : dict, optional
        Parameters for the training dataloader, see the PyTorch docs for `DataLoader`.
        If left as `None`, the dict `{"shuffle": True}` will be used, this is set in
        the `GeneralDataConfig`.
    val_dataloader_params : dict, optional
        Parameters for the validation dataloader, see PyTorch the docs for `DataLoader`.
        If left as `None`, the empty dict `{}` will be used, this is set in the
        `GeneralDataConfig`.
    checkpoint_params : dict, default=None
        Parameters for the checkpoint callback, see PyTorch Lightning documentation
        (`ModelCheckpoint`) for the list of available parameters.

    Returns
    -------
    Configuration
        Configuration for training Noise2Noise.

    Examples
    --------
    Minimum example:
    >>> config = create_n2n_configuration(
    ...     experiment_name="n2n_experiment",
    ...     data_type="array",
    ...     axes="YX",
    ...     patch_size=[64, 64],
    ...     batch_size=32,
    ...     num_epochs=100
    ... )

    You can also limit the number of batches per epoch:
    >>> config = create_n2n_configuration(
    ...     experiment_name="n2n_experiment",
    ...     data_type="array",
    ...     axes="YX",
    ...     patch_size=[64, 64],
    ...     batch_size=32,
    ...     num_steps=100  # limit to 100 batches per epoch
    ... )

    To disable transforms, simply set `augmentations` to an empty list:
    >>> config = create_n2n_configuration(
    ...     experiment_name="n2n_experiment",
    ...     data_type="array",
    ...     axes="YX",
    ...     patch_size=[64, 64],
    ...     batch_size=32,
    ...     num_epochs=100,
    ...     augmentations=[]
    ... )

    A list of transforms can be passed to the `augmentations` parameter:
    >>> from careamics.config.transformations import XYFlipModel
    >>> config = create_n2n_configuration(
    ...     experiment_name="n2n_experiment",
    ...     data_type="array",
    ...     axes="YX",
    ...     patch_size=[64, 64],
    ...     batch_size=32,
    ...     num_epochs=100,
    ...     augmentations=[
    ...         # No rotation and only Y flipping
    ...         XYFlipModel(flip_x = False, flip_y = True)
    ...     ]
    ... )

    If you are training multiple channels they will be trained independently by default,
    you simply need to specify the number of channels input (and optionally, the number
    of channels output):
    >>> config = create_n2n_configuration(
    ...     experiment_name="n2n_experiment",
    ...     data_type="array",
    ...     axes="YXC", # channels must be in the axes
    ...     patch_size=[64, 64],
    ...     batch_size=32,
    ...     num_epochs=100,
    ...     n_channels_in=3, # number of input channels
    ...     n_channels_out=1 # if applicable
    ... )

    If instead you want to train multiple channels together, you need to turn off the
    `independent_channels` parameter:
    >>> config = create_n2n_configuration(
    ...     experiment_name="n2n_experiment",
    ...     data_type="array",
    ...     axes="YXC", # channels must be in the axes
    ...     patch_size=[64, 64],
    ...     batch_size=32,
    ...     num_epochs=100,
    ...     independent_channels=False,
    ...     n_channels_in=3,
    ...     n_channels_out=1 # if applicable
    ... )

    If you would like to train on CZI files, use `"czi"` as `data_type` and `"SCYX"` as
    `axes` for 2-D or `"SCZYX"` for 3-D denoising. Note that `"SCYX"` can also be used
    for 3-D data but spatial context along the Z dimension will then not be taken into
    account.
    >>> config_2d = create_n2n_configuration(
    ...     experiment_name="n2n_experiment",
    ...     data_type="czi",
    ...     axes="SCYX",
    ...     patch_size=[64, 64],
    ...     batch_size=32,
    ...     num_epochs=100,
    ...     n_channels_in=1,
    ... )
    >>> config_3d = create_n2n_configuration(
    ...     experiment_name="n2n_experiment",
    ...     data_type="czi",
    ...     axes="SCZYX",
    ...     patch_size=[16, 64, 64],
    ...     batch_size=16,
    ...     num_epochs=100,
    ...     n_channels_in=1,
    ... )
    """
    return Configuration(
        **_create_supervised_config_dict(
            algorithm="n2n",
            experiment_name=experiment_name,
            data_type=data_type,
            axes=axes,
            patch_size=patch_size,
            batch_size=batch_size,
            trainer_params=trainer_params,
            augmentations=augmentations,
            independent_channels=independent_channels,
            loss=loss,
            n_channels_in=n_channels_in,
            n_channels_out=n_channels_out,
            logger=logger,
            model_params=model_params,
            optimizer=optimizer,
            optimizer_params=optimizer_params,
            lr_scheduler=lr_scheduler,
            lr_scheduler_params=lr_scheduler_params,
            train_dataloader_params=train_dataloader_params,
            val_dataloader_params=val_dataloader_params,
            checkpoint_params=checkpoint_params,
            num_epochs=num_epochs,
            num_steps=num_steps,
        )
    )