"""Configuration for Nemotron-Dense Audex audio-understanding HF checkpoints."""

from __future__ import annotations

from typing import Any, Optional

from .modeling_nemotron_dense import NemotronDenseConfig


class NemotronDenseAudexConfig(NemotronDenseConfig):
    """Nemotron-Dense text config plus NV-Whisper audio metadata.

    This class intentionally preserves all LLM fields from the baseline
    `NemotronDenseConfig` so existing `model.*` and `lm_head.*` weights load
    unchanged. Audio-specific fields describe the extra modules added by
    `modeling_nemotron_h_audio.py`.
    """

    model_type = "nemotron_dense_audex"
    keys_to_ignore_at_inference = ["past_key_values"]

    def __init__(
        self,
        audio_config: Optional[dict[str, Any]] = None,
        audio_model_type: str = "NV-Whisper",
        sound_model_type: Optional[str] = None,
        audio_preprocessor_path: str = "audio_preprocessor",
        sound_token: str = "<so_embedding>",
        sound_start_token: str = "<so_start>",
        sound_end_token: str = "<so_end>",
        sound_token_id: Optional[int] = None,
        sound_start_token_id: Optional[int] = None,
        sound_end_token_id: Optional[int] = None,
        sound_embedding_size: int = 750,
        sound_clip_duration: float = 30.0,
        sound_target_rate: int = 16000,
        audio_encoder_hidden_size: int = 1280,
        audio_projector_intermediate_size: int = 4096,
        audio_projector_activation: str = "relu2",
        audio_projector_norm_eps: float = 1e-5,
        **kwargs,
    ):
        self.audio_config = audio_config or {
            "model_type": "qwen2_audio_encoder",
            "num_mel_bins": 128,
            "encoder_layers": 32,
            "encoder_attention_heads": 20,
            "encoder_ffn_dim": 5120,
            "d_model": audio_encoder_hidden_size,
            "activation_function": "gelu",
            "scale_embedding": False,
            "max_source_positions": 1500,
        }
        self.audio_model_type = audio_model_type
        self.sound_model_type = sound_model_type
        self.audio_preprocessor_path = audio_preprocessor_path
        self.sound_token = sound_token
        self.sound_start_token = sound_start_token
        self.sound_end_token = sound_end_token
        self.sound_token_id = sound_token_id
        self.sound_start_token_id = sound_start_token_id
        self.sound_end_token_id = sound_end_token_id
        self.sound_embedding_size = sound_embedding_size
        self.sound_clip_duration = sound_clip_duration
        self.sound_target_rate = sound_target_rate
        self.audio_encoder_hidden_size = audio_encoder_hidden_size
        self.audio_projector_intermediate_size = audio_projector_intermediate_size
        self.audio_projector_activation = audio_projector_activation
        self.audio_projector_norm_eps = audio_projector_norm_eps
        super().__init__(**kwargs)