"""
Configuration for CAL-MOS Hugging Face models.

CAL-MOS: Bridging Layers with Adapters for Robust MOS Prediction Across
Speech Foundation Models, Interspeech 2026.
"""

from transformers import PretrainedConfig


class CalMOSConfig(PretrainedConfig):
    """Configuration for the CAL-MOS Adapters + Mean (A+M) model."""

    model_type = "calmos"

    def __init__(
        self,
        backbone_name=None,
        num_feature_layers=None,
        adapter_input_dim=None,
        adapter_dim=None,
        mlp_input_dim=None,
        mlp_hidden_dim=None,
        mlp_num_layers=1,
        mlp_output_size=1,
        mlp_dropout=0.1,
        target_sr=16000,
        mos_min=1.0,
        mos_max=5.0,
        layer_weight_strategy="layer_adapter",
        pooling_strategy="mean",
        freeze_backbone=True,
        **kwargs,
    ):
        super().__init__(**kwargs)

        self.backbone_name = backbone_name
        self.num_feature_layers = num_feature_layers
        self.adapter_input_dim = adapter_input_dim
        self.adapter_dim = adapter_dim
        self.mlp_input_dim = mlp_input_dim
        self.mlp_hidden_dim = mlp_hidden_dim
        self.mlp_num_layers = mlp_num_layers
        self.mlp_output_size = mlp_output_size
        self.mlp_dropout = mlp_dropout

        self.target_sr = target_sr
        self.mos_min = mos_min
        self.mos_max = mos_max

        self.layer_weight_strategy = layer_weight_strategy
        self.pooling_strategy = pooling_strategy
        self.freeze_backbone = freeze_backbone