#
#  Copyright (c) 2025
#  Minh NGUYEN <vnguyen9@lakeheadu.ca>
#
import logging

from typing import Literal
from transformers import PretrainedConfig

logger = logging.getLogger(__name__)


class ResnetConfig(PretrainedConfig):
    """The config class for custom resnet."""

    model_type: str = "custom-resnet"
    """The unique model type used to write meta on HuggingFace and register auto class."""

    def __init__(
        self,
        block_type: Literal['bottleneck', 'basic'] = "bottleneck",
        layers: tuple[int] = (3, 4, 6, 3),
        num_classes: int = 1000,
        input_channels: int = 3,
        cardinality: int = 1,
        base_width: int = 64,
        stem_width: int = 64,
        stem_type: Literal['', 'deep', 'deep-tiered'] = "",
        avg_down: bool = False,
        load_pretrained_state: bool = True,
        pretrained_state_source: Literal['timm'] = "timm",
        **kwargs
    ):
        if block_type not in ["basic", "bottleneck"]:
            raise ValueError(f"`block_type` must be 'basic' or 'bottleneck', got {block_type!r}.")
        if stem_type not in ["", "deep", "deep-tiered"]:
            raise ValueError(f"`stem_type` must be '', 'deep' or 'deep_tiered', got {stem_type!r}")

        # config used to create model
        self.block_type = block_type
        self.layers = layers
        self.num_classes = num_classes
        self.input_channels = input_channels
        self.cardinality = cardinality
        self.base_width = base_width
        self.stem_width = stem_width
        self.stem_type = stem_type
        self.avg_down = avg_down

        self.load_pretrained_state = load_pretrained_state
        self.pretrained_state_source = pretrained_state_source

        super().__init__(**kwargs)