"""Configuration for the Mettle pathology tile encoder."""

from transformers import PretrainedConfig


class MettleConfig(PretrainedConfig):
    """Configuration for a Mettle ViT-g/14 encoder."""

    model_type = "mettle"

    def __init__(
        self,
        backbone_name="vit_giant_patch14_reg4_dinov2",
        image_size=224,
        hidden_size=1536,
        num_register_tokens=4,
        head_enabled=True,
        head_num_atoms=16,
        head_rank=8,
        head_hidden_size=256,
        default_feature_view="cls",
        **kwargs,
    ):
        super().__init__(**kwargs)
        if default_feature_view not in {"cls", "cls_mean"}:
            raise ValueError(
                "default_feature_view must be 'cls' or 'cls_mean', got "
                f"{default_feature_view!r}"
            )
        self.backbone_name = backbone_name
        self.image_size = int(image_size)
        self.hidden_size = int(hidden_size)
        self.num_register_tokens = int(num_register_tokens)
        self.head_enabled = bool(head_enabled)
        self.head_num_atoms = int(head_num_atoms)
        self.head_rank = int(head_rank)
        self.head_hidden_size = int(head_hidden_size)
        self.default_feature_view = default_feature_view