from typing import Optional, Union
from transformers import Qwen2Config
from transformers.configuration_utils import PretrainedConfig


class UASAudioEncoderConfig(PretrainedConfig):
    model_type = "uas_audio_encoder"
    def __init__(
        self,
        num_mel_bins: Optional[int] = 128,
        encoder_layers: Optional[int] = 32,
        encoder_attention_heads: Optional[int] = 20,
        encoder_ffn_dim: Optional[int] = 5120,
        d_model: Optional[int] = 1280,
        dropout: Optional[int] = 0,
        attention_dropout: Optional[int] = 0,
        activation_function: Optional[int] = "gelu",
        activation_dropout: Optional[int] = 0,
        initializer_range: Optional[int] = 0.02,
        max_source_positions: Optional[int] = 1500,
        n_window: Optional[int] = 50,
        n_window_infer: Optional[int] = 800,
        conv_chunksize: Optional[int] = 500,
        downsample_hidden_size: Optional[int] = 480,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.num_mel_bins = num_mel_bins
        self.d_model = d_model
        self.encoder_layers = encoder_layers
        self.encoder_attention_heads = encoder_attention_heads
        self.encoder_ffn_dim = encoder_ffn_dim
        self.dropout = dropout
        self.attention_dropout = attention_dropout
        self.activation_function = activation_function
        self.activation_dropout = activation_dropout
        self.num_hidden_layers = encoder_layers
        self.initializer_range = initializer_range
        self.max_source_positions = max_source_positions
        self.n_window = n_window
        self.n_window_infer = n_window_infer
        self.conv_chunksize = conv_chunksize
        self.downsample_hidden_size = downsample_hidden_size


class UASAudioTextConfig(PretrainedConfig):
    model_type = "uas_audio_text"

    def __init__(
        self,
        vocab_size=64012,
        hidden_size=4096,
        intermediate_size=11008,
        num_hidden_layers=48,
        num_attention_heads=32,
        num_attention_groups=4,
        num_key_value_heads=4,
        hidden_act="silu",
        max_position_embeddings=8192,
        initializer_range=0.02,
        rms_norm_eps=1e-6,
        rope_theta=1000000.0,
        rope_scaling=None,
        eos_token_id=None,
        **kwargs
    ):

        super().__init__(
            **kwargs)

        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.num_attention_groups = num_attention_groups
        self.num_key_value_heads = num_key_value_heads
        assert self.num_attention_groups == self.num_key_value_heads, \
            "num_attention_groups must be equal to num_key_value_heads"
        self.hidden_act = hidden_act
        self.max_position_embeddings = max_position_embeddings
        self.initializer_range = initializer_range
        self.rms_norm_eps = rms_norm_eps
        self.rope_theta = rope_theta
        self.rope_scaling = rope_scaling
        self.eos_token_id = eos_token_id

        self.text_config = Qwen2Config(
            vocab_size=vocab_size,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_hidden_layers=num_hidden_layers,
            num_attention_heads=num_attention_heads,
            num_key_value_heads=num_key_value_heads,
            hidden_act=hidden_act,
            max_position_embeddings=max_position_embeddings,
            initializer_range=initializer_range,
            rms_norm_eps=rms_norm_eps,
            rope_theta=rope_theta,
            rope_scaling=rope_scaling,
            architectures=["Qwen2ForCausalLM"],
            dtype=getattr(self, "dtype", "bfloat16"),
        )


class UASAudioConfig(PretrainedConfig):
    model_type = "uas_audio"
    architectures = ["UASAudioForCausalLM"]

    def __init__(
        self,
        audio_encoder_config: Optional[Union[dict, UASAudioEncoderConfig]] = None,
        text_config: Optional[Union[dict, UASAudioTextConfig]] = None,
        **kwargs
    ):
        super().__init__(**kwargs)
        if text_config is None:
            text_config = UASAudioTextConfig().text_config
        elif isinstance(text_config, dict):
            text_config = UASAudioTextConfig(**text_config).text_config
        self.text_config = text_config

        if audio_encoder_config is None:
            self.audio_encoder_config = UASAudioEncoderConfig()
        elif isinstance(audio_encoder_config, dict):
            self.audio_encoder_config = UASAudioEncoderConfig(**audio_encoder_config)
        elif isinstance(audio_encoder_config, UASAudioEncoderConfig):
            self.audio_encoder_config = audio_encoder_config


class UASAudioEncoderOnlyConfig(PretrainedConfig):
    model_type = "uas_audio_encoder_only"
    architectures = ["UASAudioEncoderOnly"]

    def __init__(
        self,
        audio_encoder_config: Optional[Union[dict, UASAudioEncoderConfig]] = None,
        hidden_size: Optional[int] = 4096,  # LLM hidden size for adapter output
        **kwargs
    ):
        super().__init__(**kwargs)
        if audio_encoder_config is None:
            self.audio_encoder_config = UASAudioEncoderConfig()
        elif isinstance(audio_encoder_config, dict):
            self.audio_encoder_config = UASAudioEncoderConfig(**audio_encoder_config)
        elif isinstance(audio_encoder_config, UASAudioEncoderConfig):
            self.audio_encoder_config = audio_encoder_config
        self.hidden_size = hidden_size  # Output dimension of adapter