# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Configuration for Backbone models"""

from huggingface_hub.dataclasses import strict

from ...backbone_utils import BackboneConfigMixin
from ...configuration_utils import PreTrainedConfig
from ...utils import auto_docstring


@auto_docstring(checkpoint="")
@strict
class TimmBackboneConfig(BackboneConfigMixin, PreTrainedConfig):
    r"""
    backbone (`str`, *optional*):
        The timm checkpoint to load.
    features_only (`bool`, *optional*, defaults to `True`):
        Whether to output only the features or also the logits.
    freeze_batch_norm_2d (`bool`, *optional*, defaults to `False`):
        Converts all `BatchNorm2d` and `SyncBatchNorm` layers of provided module into `FrozenBatchNorm2d`.

    Example:
    ```python
    >>> from transformers import TimmBackboneConfig, TimmBackbone

    >>> # Initializing a timm backbone
    >>> configuration = TimmBackboneConfig("resnet50")

    >>> # Initializing a model from the configuration
    >>> model = TimmBackbone(configuration)

    >>> # Accessing the model configuration
    >>> configuration = model.config
    ```
    """

    model_type = "timm_backbone"

    backbone: str | None = None
    num_channels: int = 3
    features_only: bool = True
    _out_indices: list[int] | None = None
    freeze_batch_norm_2d: bool = False
    output_stride: int | None = None

    def __post_init__(self, **kwargs):
        self.out_indices = self.out_indices if self.out_indices is not None else [-1]
        super().__post_init__(**kwargs)

    @property
    def out_indices(self):
        return self._out_indices

    @out_indices.setter
    def out_indices(self, out_indices: tuple[int, ...] | list[int]):
        """
        Set the out_indices attribute. This will also update the out_features attribute to match the new out_indices.
        """
        self._out_indices = list(out_indices) if out_indices is not None else out_indices
        if getattr(self, "stage_names", None) is not None:
            self.set_output_features_output_indices(out_features=None, out_indices=out_indices)

    @property
    def out_features(self):
        return self._out_features

    @out_features.setter
    def out_features(self, out_features: list[str]):
        """
        Set the out_features attribute. This will also update the out_indices attribute to match the new out_features.
        """
        self._out_features = out_features
        if getattr(self, "stage_names", None) is not None:
            self.set_output_features_output_indices(out_features=out_features, out_indices=None)


__all__ = ["TimmBackboneConfig"]
