"""MERaLiON2 model configuration"""

from transformers import Gemma2Config, WhisperConfig
from transformers.configuration_utils import PretrainedConfig
from transformers.utils import logging


logger = logging.get_logger(__name__)


class MERaLiON3Config(PretrainedConfig):
    r"""
    This is the configuration class to store the configuration of a [`MERaLiON3ForConditionalGeneration`]. It is used to instantiate an
    MERaLiON3 model according to the specified arguments, defining the model architecture. Instantiating a configuration
    with the defaults will yield a similar configuration to that of the MERaLiON3.

    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
    documentation from [`PretrainedConfig`] for more information.

    Args:
        audio_config (`Union[AutoConfig, dict]`,  *optional*, defaults to `CLIPVisionConfig`):
            The config object or dictionary of the audio backbone.
        text_config (`Union[AutoConfig, dict]`, *optional*, defaults to `LlamaConfig`):
            The config object or dictionary of the text backbone.
        audio_token_index (`int`, *optional*, defaults to 151646):
            The image token index to encode the image prompt.
    """

    model_type = "meralion3"
    is_composition = False

    def __init__(
        self,
        speech_config=None,
        text_config=None,
        speech_mlp_use_projection=True,
        speech_mlp_scale_factor=5,
        speech_token_index=255999,
        fixed_speech_embeds_length=None,
        **kwargs,
    ):
        # Set as instance attribute so it's serialized into config.json,
        # enabling trust_remote_code=True loading without manual registration.
        self.auto_map = {
            "AutoConfig": "configuration_meralion3.MERaLiON3Config",
            "AutoModelForSpeechSeq2Seq": "modeling_meralion3.MERaLiON3ForConditionalGeneration",
            "AutoProcessor": "processing_meralion3.MERaLiON3Processor",
        }

        if isinstance(speech_config, dict):
            speech_config = WhisperConfig(**speech_config)
        elif speech_config is None:
            speech_config = WhisperConfig(
                d_model=1280,
                encoder_attention_heads=20,
                encoder_ffn_dim=5120,
                encoder_layerdrop=0.0,
                encoder_layers=32,
                num_mel_bins=128,
                max_source_positions=1500,
                scale_embedding=False,
                activation_function="gelu",
            )

        self.speech_config = speech_config

        if isinstance(text_config, dict):
            text_config = Gemma2Config(**text_config)
        elif text_config is None:
            text_config = Gemma2Config()

        self.text_config = text_config

        self.speech_mlp_use_projection = speech_mlp_use_projection
        self.speech_mlp_scale_factor = speech_mlp_scale_factor
        self.speech_token_index = speech_token_index
        # Whisper encoder outputs 1500 tokens (3000 mel frames / stride 2).
        # The adapter reduces by speech_mlp_scale_factor, so the number of
        # speech embedding tokens per chunk is 1500 // scale_factor.
        if fixed_speech_embeds_length is None:
            fixed_speech_embeds_length = 1500 // speech_mlp_scale_factor
        self.fixed_speech_embeds_length = fixed_speech_embeds_length
        
        self.sliding_window = self.text_config.sliding_window
        self.hidden_size = self.text_config.hidden_size
        self.num_attention_heads = self.text_config.num_attention_heads
        self.num_hidden_layers = self.text_config.num_hidden_layers
        self.num_key_value_heads = self.text_config.num_key_value_heads
        self.head_dim = self.text_config.head_dim
        self.intermediate_size = self.text_config.intermediate_size
        # Gemma2 requires HybridCache for correct sliding window attention
        # on alternating layers. Without this, GenerationMixin defaults to
        # DynamicCache which disables sliding window and causes repetition.
        self.cache_implementation = getattr(
            self.text_config, 'cache_implementation', 'hybrid'
        )
        
        super().__init__(**kwargs)