"""Configuration for the standalone RepConTUR HF model."""

from transformers import PretrainedConfig

from ._utils import PHI3_CHAT_TEMPLATE


class RepConTURConfig(PretrainedConfig):
    """Config for RepConTUR: InternVL backbone + CT-ViT + VoxTell 3D grounding.

    The InternVL backbone architecture is pulled from ``internvl_model_path``
    via its hub remote code at load time; all RepConTUR weights (backbone,
    CT encoder, pooler, VoxTell, bridge) ship in the model's safetensors.
    """

    model_type = 'repcontur'

    def __init__(
        self,
        internvl_model_path='OpenGVLab/InternVL2_5-1B',
        vocab_size=151679,
        llm_hidden_size=896,
        ct_encoder=None,
        ct_pooler=None,
        voxtell_arch_kwargs=None,
        voxtell_kw_requires_import=('conv_op', 'norm_op', 'dropout_op', 'nonlin'),
        voxtell_decoder_layer=4,
        voxtell_text_embedding_dim=2560,
        voxtell_num_maskformer_stages=5,
        voxtell_num_heads=32,
        voxtell_query_dim=2048,
        special_tokens=None,
        seg_token_idx=None,
        img_context_token_id=None,
        template=None,
        generation_kwargs=None,
        torch_dtype='bfloat16',
        **kwargs,
    ):
        self.internvl_model_path = internvl_model_path
        self.vocab_size = vocab_size
        self.llm_hidden_size = llm_hidden_size
        self.ct_encoder = ct_encoder or {}
        self.ct_pooler = ct_pooler or {}
        self.voxtell_arch_kwargs = voxtell_arch_kwargs or {}
        self.voxtell_kw_requires_import = list(voxtell_kw_requires_import)
        self.voxtell_decoder_layer = voxtell_decoder_layer
        self.voxtell_text_embedding_dim = voxtell_text_embedding_dim
        self.voxtell_num_maskformer_stages = voxtell_num_maskformer_stages
        self.voxtell_num_heads = voxtell_num_heads
        self.voxtell_query_dim = voxtell_query_dim
        self.special_tokens = special_tokens or ['[SEG]', '<p>', '</p>', '<vp>', '</vp>']
        self.seg_token_idx = seg_token_idx
        self.img_context_token_id = img_context_token_id
        self.template = template or dict(PHI3_CHAT_TEMPLATE)
        self.generation_kwargs = generation_kwargs or {}
        super().__init__(torch_dtype=torch_dtype, **kwargs)