"""
MATPAC model configuration
"""

from transformers.configuration_utils import PretrainedConfig
from transformers.utils import logging

logger = logging.get_logger(__name__)


class MatpacConfig(PretrainedConfig):
    r"""
    Configuration for [`MatpacModel`]. Every checkpoint released for MATPAC / MATPAC++
    (https://github.com/aurianworld/matpac) maps to one set of these values:

    | checkpoint                                    | n_t | as_class_head | concat_freq | probe_out_features |
    |------------------------------------------------|-----|----------------|-------------|--------------------|
    | matpac_10_2048.pt                               | 608 | False          | True        | None               |
    | matpac_plus_6s_2048_enconly.pt                  | 608 | False          | True        | None               |
    | matpac_plus_music_6s_2048_enconly.pt            | 608 | False          | True        | None               |
    | matpac_plus_as_48_1_map_enconly.pt              | 992 | False          | True        | None               |
    | matpac_plus_as_48_1_map_enc_and_head.pt         | 992 | True           | False       | None               |
    | matpac_plus_6s_2048_enconly.pt + probe/<task>   | 608 | False          | True        | task dependent     |
    | matpac_plus_music_6s_2048_enconly.pt + probe/<task> | 608 | False      | True        | task dependent     |
    """

    model_type = "matpac"

    def __init__(
        self,
        embed_dim=768,
        depth=12,
        num_heads=12,
        mlp_ratio=4,
        n_freq=80,
        n_t=608,
        patch_size=16,
        sample_rate=16000,
        lms_mean=-7.056,
        lms_std=4.193,
        inference_type="precise",
        pull_time_dimension=True,
        concat_freq=True,
        as_class_head=False,
        probe_out_features=None,
        **kwargs,
    ):
        super().__init__(**kwargs)

        # Encoder (ViT) parameters
        self.embed_dim = embed_dim
        self.depth = depth
        self.num_heads = num_heads
        self.mlp_ratio = mlp_ratio

        # Log-mel spectrogram / patchification parameters
        self.n_freq = n_freq
        self.n_t = n_t
        self.patch_size = patch_size
        self.sample_rate = sample_rate
        self.lms_mean = lms_mean
        self.lms_std = lms_std

        # Inference behaviour
        self.inference_type = inference_type  # "precise" or "fast"
        self.pull_time_dimension = pull_time_dimension
        self.concat_freq = concat_freq

        # Optional heads. `as_class_head` is the AudioSet fine-tuning head
        # (uses `self.id2label`/`self.num_labels`, 527 classes). `probe_out_features`
        # instantiates a downstream linear probe (see `class_mappings/` in the
        # original repo for the label sets, now expected as `id2label` in `kwargs`).
        self.as_class_head = as_class_head
        self.probe_out_features = probe_out_features