# -*- coding: utf-8 -*-
"""
Configuration class for MiniCPM5-V (Multimodal Vision & Document & Video on MiniCPM5-2B).
"""
import copy
from transformers.configuration_utils import PretrainedConfig
from transformers.models.llama.configuration_llama import LlamaConfig
from transformers.models.siglip.configuration_siglip import SiglipVisionConfig


class MiniCPM5VSliceConfig(PretrainedConfig):
    r"""Configuration for dynamic image slicing in MiniCPM5-V."""
    def __init__(
        self,
        max_slice_nums=9,
        scale_resolution=448,
        patch_size=14,
        slice_mode=True,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.max_slice_nums = max_slice_nums
        self.scale_resolution = scale_resolution
        self.patch_size = patch_size
        self.slice_mode = slice_mode


class MiniCPM5VConfig(PretrainedConfig):
    r"""
    Unified configuration class for MiniCPM5-V combining:
    - Language Backbone: MiniCPM5-2B (LlamaConfig, hidden_size=2048, layers=42, vocab=130560)
    - Vision Backbone: SigLIP (SiglipVisionConfig, hidden_size=1152, patch_size=14)
    - Resampler: Spatial 2x2 Pixel Unshuffle (4608 -> 2048 -> 2048)
    - High-Res Dynamic Slicing: up to 9 slices (1.8M pixels) + overview
    """
    model_type = "minicpm5_v"
    is_composition = True

    def __init__(
        self,
        text_config=None,
        vision_config=None,
        slice_config=None,
        projector_hidden_act="gelu",
        spatial_downsample_factor=2,
        drop_vision_last_layer=False,
        image_token_id=130074,
        slice_start_token_id=130075,
        slice_end_token_id=130076,
        box_start_token_id=130077,
        box_end_token_id=130078,
        ref_start_token_id=130079,
        ref_end_token_id=130080,
        **kwargs
    ):
        super().__init__(**kwargs)

        if text_config is None:
            # Default to official openbmb/MiniCPM5-2B specifications
            text_config = {
                "vocab_size": 130560,
                "hidden_size": 2048,
                "intermediate_size": 6144,
                "num_hidden_layers": 42,
                "num_attention_heads": 16,
                "num_key_value_heads": 2,
                "head_dim": 128,
                "max_position_embeddings": 131072,
                "rms_norm_eps": 1e-6,
                "rope_theta": 5000000.0,
                "hidden_act": "silu",
                "tie_word_embeddings": False,
                "torch_dtype": "bfloat16",
            }
        
        if vision_config is None:
            # Default to OpenBMB SigLIP-so400m-980/448 visual specifications
            vision_config = {
                "hidden_size": 1152,
                "image_size": 448,
                "intermediate_size": 4304,
                "num_attention_heads": 16,
                "num_hidden_layers": 27,
                "patch_size": 14,
                "hidden_act": "gelu_pytorch_tanh",
                "layer_norm_eps": 1e-6,
            }

        if slice_config is None:
            slice_config = {
                "max_slice_nums": 9,
                "scale_resolution": 448,
                "patch_size": 14,
                "slice_mode": True,
            }

        if isinstance(text_config, dict):
            self.text_config = LlamaConfig(**text_config)
        else:
            self.text_config = text_config

        if isinstance(vision_config, dict):
            self.vision_config = SiglipVisionConfig(**vision_config)
        else:
            self.vision_config = vision_config

        if isinstance(slice_config, dict):
            self.slice_config = MiniCPM5VSliceConfig(**slice_config)
        else:
            self.slice_config = slice_config

        self.projector_hidden_act = projector_hidden_act
        self.spatial_downsample_factor = spatial_downsample_factor
        self.drop_vision_last_layer = drop_vision_last_layer

        # Special Token IDs for OCR & Dynamic Slicing
        self.image_token_id = image_token_id
        self.slice_start_token_id = slice_start_token_id
        self.slice_end_token_id = slice_end_token_id
        self.box_start_token_id = box_start_token_id
        self.box_end_token_id = box_end_token_id
        self.ref_start_token_id = ref_start_token_id
        self.ref_end_token_id = ref_end_token_id

    def to_dict(self):
        output = copy.deepcopy(self.__dict__)
        output["text_config"] = self.text_config.to_dict()
        output["vision_config"] = self.vision_config.to_dict()
        output["slice_config"] = self.slice_config.to_dict()
        output["model_type"] = self.__class__.model_type
        return output