"""Configuration class for the CoralBay Swin UNETR-based encoder model."""

from typing import Literal

from transformers import PretrainedConfig


class CoralBayConfig(PretrainedConfig):
    """Configuration for CoralBay, a Swin UNETR-based encoder model.

    Args:
        in_channels: Number of input image channels.
        feature_size: Base feature dimension for the Swin Transformer.
        spatial_dims: Number of spatial dimensions (2 or 3).
        out_indices: Number of feature outputs. If None,
            the aggregated feature vector is returned.
        embeddings_type: Output embedding style — ``"multiscale"`` aggregates
            all encoder stages; ``"head"`` uses only the last stage.
        use_v2: Whether to use SwinTransformerV2 weights and architecture.
    """

    model_type = "coralbay"

    def __init__(
        self,
        in_channels: int = 1,
        feature_size: int = 48,
        spatial_dims: int = 3,
        out_indices: int | None = None,
        embeddings_type: Literal["multiscale", "head"] = "multiscale",
        use_v2: bool = True,
        **kwargs
    ) -> None:
        super().__init__(**kwargs)

        self.in_channels = in_channels
        self.feature_size = feature_size
        self.spatial_dims = spatial_dims
        self.out_indices = out_indices
        self.embeddings_type = embeddings_type
        self.use_v2 = use_v2