from __future__ import annotations

from transformers import PretrainedConfig


class MiniCPMVLAConfig(PretrainedConfig):
    """Configuration for the complete VLM and action-head VLA model."""

    model_type = "minicpm_vla"
    is_composition = True

    def __init__(
        self,
        vlm_config: dict | PretrainedConfig | None = None,
        action_dim: int = 80,
        state_dim: int = 80,
        action_horizon: int = 30,
        max_num_embodiments: int = 32,
        num_inference_timesteps: int = 4,
        vlm_dtype: str = "bfloat16",
        action_head_dtype: str = "float32",
        **kwargs,
    ):
        super().__init__(**kwargs)
        if isinstance(vlm_config, PretrainedConfig):
            vlm_config = vlm_config.to_dict()
        self.vlm_config = vlm_config
        self.action_dim = action_dim
        self.state_dim = state_dim
        self.action_horizon = action_horizon
        self.max_num_embodiments = max_num_embodiments
        self.num_inference_timesteps = num_inference_timesteps
        self.vlm_dtype = vlm_dtype
        self.action_head_dtype = action_head_dtype
        self.architectures = ["MiniCPMV_VLA"]