"""Configuration for the self-contained Jolia Atlas backbone."""

from __future__ import annotations

from transformers import PretrainedConfig


class JoliaConfig(PretrainedConfig):
    """Config for :class:`~modeling_jolia.JoliaModel`.

    Jolia wraps a ``MultiModalAtlas`` 3D vision backbone and, when the source
    run was trained with a ParallelOrganCLIP objective, a per-scale bank of
    organ-query cross-attention pooling heads.

    Args:
        embed_dim: Pooled (CLS-equivalent) embedding size.
        multiscale_feats: Whether the backbone concatenates multi-scale pooled
            features (always ``True`` for the released checkpoint).
        atlas_config: Nested dict describing the Atlas modalities / stages.
            Passed straight through to ``MultiModalAtlasConfig``.
        patch_embed_dim: Per-scale organ-query dimension (``dim`` of the Atlas
            stages). Total ``query_dim`` is ``num_scales * patch_embed_dim``.
        num_organs: Number of learnable organ-query slots per scale.
        num_scales: Number of multi-scale levels with an organ-query head.
        num_heads: Attention heads in each organ-query pooler.
        use_cls_residual: Add the CLS vector to every organ query (matches the
            training-time ParallelOrganCLIP setting).
        organ_slot_names: Ordered names for the organ-query slots. Length may be
            ``< num_organs`` — trailing slots are unused padding. Enables the
            named :meth:`JoliaModel.encode_organs` API.
        text_embed_dim: Hidden size of the paired text encoder. Set non-zero to
            build the linear text-projection head used for CLIP-style zero-shot
            (maps text features into the same ``embed_dim`` space as the
            vision CLS). ``0`` -> no text head.
        text_encoder_id: HuggingFace id of the text encoder Jolia was trained
            against (e.g. ``"Qwen/Qwen3-Embedding-8B"``). Informational; loading
            it is opt-in via :class:`text_encoder_jolia.JoliaTextEncoder`.
        text_context_length: Tokenizer ``max_length`` to match training.
    """

    model_type = "jolia"

    def __init__(
        self,
        embed_dim: int = 576,
        multiscale_feats: bool = True,
        atlas_config: dict | None = None,
        patch_embed_dim: int = 0,
        num_organs: int = 0,
        num_scales: int = 0,
        num_heads: int = 1,
        use_cls_residual: bool = False,
        organ_slot_names: list[str] | None = None,
        text_embed_dim: int = 0,
        text_encoder_id: str = "",
        text_context_length: int = 512,
        **kwargs: object,
    ) -> None:
        self.embed_dim = embed_dim
        self.multiscale_feats = multiscale_feats
        self.atlas_config = atlas_config or {}
        self.patch_embed_dim = patch_embed_dim
        self.num_organs = num_organs
        self.num_scales = num_scales
        self.num_heads = num_heads
        self.use_cls_residual = use_cls_residual
        self.organ_slot_names = list(organ_slot_names) if organ_slot_names else []
        self.text_embed_dim = text_embed_dim
        self.text_encoder_id = text_encoder_id
        self.text_context_length = text_context_length
        super().__init__(**kwargs)

    @property
    def has_text_head(self) -> bool:
        return self.text_embed_dim > 0

    @property
    def query_dim(self) -> int:
        """Total organ-query dimension (concatenated across scales)."""
        return int(self.num_scales * self.patch_embed_dim)

    @property
    def has_queries(self) -> bool:
        return self.num_organs > 0 and self.num_scales > 0