from typing import Optional
from transformers import ModernBertConfig, DacConfig


DFLT_DACVAE_CONFIG = {
    "encoder_hidden_size": 64,
    "downsampling_ratios": [2, 8, 10, 12],
    "decoder_hidden_size": 1536,
    "n_codebooks": 16,
    "codebook_size": 1024,
    "codebook_dim": 128,
    "quantizer_dropout": 0,
    "sampling_rate": 48000,
}

DFLT_TEXT_ENCODER_CONFIG = {
    "classifier_pooling": "mean",
    "global_attn_every_n_layers": 3,
    "global_rope_theta": 160000.0,
    "hidden_size": 1024,
    "intermediate_size": 2624,
    "layer_norm_eps": 1e-5,
    "local_rope_theta": 10000.0,
    "model_type": "modernbert",
    "num_attention_heads": 16,
    "num_hidden_layers": 28,
    "position_embedding_type": "absolute",
    "tie_word_embeddings": True,
    "torch_dtype": "float32",
}


class TransformerConfig:
    def __init__(
        self,
        hidden_size=1024,
        intermediate_size=2752,
        num_hidden_layers=16,
        num_attention_heads=8,
        num_key_value_heads=None,
        hidden_act="silu",
        max_position_embeddings=10_000,
        rms_norm_eps=1e-5,
        rope_theta=20000.0,
        rope_scaling=None,
        attention_bias=False,
        attention_dropout=0.0,
        **kwargs,
    ):
        self.max_position_embeddings = max_position_embeddings
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        if num_key_value_heads is None:
            num_key_value_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.head_dim = hidden_size // num_attention_heads
        self.hidden_act = hidden_act
        self.rms_norm_eps = rms_norm_eps
        self.rope_theta = rope_theta
        self.rope_scaling = rope_scaling
        self.attention_bias = attention_bias
        self.attention_dropout = attention_dropout


class AudioEncoderConfig:
    def __init__(
        self,
        dac_vae_encoder: Optional[dict] = None,
        audio_transformer: Optional[dict] = None,
        **kwargs,
    ):
        dac_vae_encoder = dac_vae_encoder or DFLT_DACVAE_CONFIG
        audio_transformer = audio_transformer or {}
        self.dac_vae_encoder = DacConfig(**dac_vae_encoder)
        self.audio_transformer = TransformerConfig(**audio_transformer)


class VisualEncoderConfig:
    def __init__(
        self,
        pe_encoder: str = "PE-Core-L14-336",
        visual_transformer: Optional[dict] = None,
        fixed_len_video: bool = False,
        **kwargs,
    ):
        visual_transformer = visual_transformer or {}

        self.pe_encoder = pe_encoder
        self.visual_transformer = TransformerConfig(**visual_transformer)
        self.fixed_len_video = fixed_len_video


class PEAudioVisualEncoderConfig:
    def __init__(
        self,
        audio_visual_transformer: Optional[dict] = None,
        visual_model: Optional[dict] = None,
        audio_model: Optional[dict] = None,
        **kwargs,
    ):
        visual_model = visual_model or {}
        audio_visual_transformer = audio_visual_transformer or {}
        audio_model = audio_model or {}

        self.visual_model = VisualEncoderConfig(**visual_model)
        self.audio_model = AudioEncoderConfig(**audio_model)
        self.audio_visual_transformer = TransformerConfig(**audio_visual_transformer)


class PEAudioVisualConfig:
    def __init__(
        self,
        audio_visual_model: Optional[dict] = None,
        text_model: Optional[dict] = None,
        output_dim: int = 1024,
        nth_text_layer: Optional[int] = 22,
        **kwargs,
    ):
        text_model = text_model or DFLT_TEXT_ENCODER_CONFIG
        audio_visual_model = audio_visual_model or {}
        self.text_model = ModernBertConfig(**text_model)
        self.audio_visual_model = PEAudioVisualEncoderConfig(**audio_visual_model)
        self.output_dim = output_dim
        self.nth_text_layer = nth_text_layer


class PEAudioFrameConfig:
    def __init__(
        self,
        audio_model: Optional[dict] = None,
        text_model: Optional[dict] = None,
        output_dim: int = 1024,
        nth_text_layer: Optional[int] = 22,
        **kwargs,
    ):
        text_model = text_model or DFLT_TEXT_ENCODER_CONFIG
        audio_model = audio_model or {}
        self.text_model = ModernBertConfig(**text_model)
        self.audio_model = AudioEncoderConfig(**audio_model)
        self.output_dim = output_dim
        self.nth_text_layer = nth_text_layer