from typing import Literal

from transformers import PretrainedConfig
from transformers.utils import logging

from .configuration_msp_audio import MSPAudioConfig
from .configuration_msp_visual import MSPVisualConfig
from .configuration_msp_fusion import MSPFusionConfig

logger = logging.get_logger(__name__)


class MSPConfig(PretrainedConfig):
    model_type = "msp"

    sub_configs = {
        "audio_config": MSPAudioConfig,
        "visual_config": MSPVisualConfig,
        "msp_fusion_config": MSPFusionConfig,
    }

    def __init__(
        self,
        audio_config: MSPAudioConfig | dict | None = None,
        visual_config: MSPVisualConfig | dict | None = None,
        msp_fusion_config: MSPFusionConfig | dict | None = None,
        final_dropout: float = 0.1,
        vocab_size: int = 32,
        ctc_loss_reduction: Literal["mean", "sum", "none"] = "mean",
        ctc_zero_infinity: bool = True,
        ctc_loss_audio_weight: float = 0.0,
        ctc_loss_visual_weight: float = 0.0,
        ctc_loss_msp_weight: float = 1.0,
        modality_dropout_prob: float = 0.0,
        audio_dropout_prob: float = 0.5,
        visual_dropout_prob: float = 0.5,
        pad_token_id: int = 0,
        bos_token_id: int = 1,
        eos_token_id: int = 2,
        **kwargs,
    ):
        # compatibility مع config.json القديم
        if msp_fusion_config is None and "fusion_config" in kwargs:
            msp_fusion_config = kwargs.pop("fusion_config")

        super().__init__(**kwargs)

        if isinstance(audio_config, dict):
            audio_config = MSPAudioConfig(**audio_config)
        elif audio_config is None:
            audio_config = MSPAudioConfig()

        if isinstance(visual_config, dict):
            visual_config = MSPVisualConfig(**visual_config)
        elif visual_config is None:
            visual_config = MSPVisualConfig()

        if isinstance(msp_fusion_config, dict):
            msp_fusion_config = MSPFusionConfig(**msp_fusion_config)
        elif msp_fusion_config is None:
            msp_fusion_config = MSPFusionConfig(
                audio_hidden_size=audio_config.hidden_size,
                visual_hidden_size=visual_config.hidden_size,
                fusion_hidden_size=max(
                    audio_config.hidden_size, visual_config.hidden_size
                ),
            )

        self.audio_config = audio_config
        self.visual_config = visual_config
        self.msp_fusion_config = msp_fusion_config

        self.final_dropout = final_dropout
        self.vocab_size = vocab_size
        self.ctc_loss_reduction = ctc_loss_reduction
        self.ctc_zero_infinity = ctc_zero_infinity
        self.ctc_loss_audio_weight = ctc_loss_audio_weight
        self.ctc_loss_visual_weight = ctc_loss_visual_weight
        self.ctc_loss_msp_weight = ctc_loss_msp_weight
        self.modality_dropout_prob = modality_dropout_prob
        self.audio_dropout_prob = audio_dropout_prob
        self.visual_dropout_prob = visual_dropout_prob
        self.pad_token_id = pad_token_id
        self.bos_token_id = bos_token_id
        self.eos_token_id = eos_token_id