# coding=utf-8
# Copyright 2024 NAVER Cloud Corp. and 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.
"""TATok (Text-Aligned Tokenizer) configuration"""

import copy
from typing import Optional

from transformers import AutoConfig
from transformers.configuration_utils import PretrainedConfig


class TATokConfig(PretrainedConfig):
    r"""
    This is the configuration class to store the configuration of a TATok (Text-Aligned Tokenizer) model.

    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
    documentation from [`PretrainedConfig`] for more information.

    Args:
        bottleneck (`dict`, *optional*):
            Configuration dict for the bottleneck layer. If not provided, uses `TATokConfig._default_bottleneck`.
        bottleneck_token_num (`int`, *optional*, defaults to 729):
            Number of bottleneck tokens.
        input_size (`int`, *optional*, defaults to 384):
            Input image size.
        teacher (`str`, *optional*, defaults to `"google/siglip2-so400m-patch14-384"`):
            Name or path of the teacher model.
        input_type (`str`, *optional*, defaults to `"indices"`):
            Type of input representation.
        pool_scale (`int`, *optional*, defaults to 1):
            Pooling scale factor.
        decoder_depth (`int`, *optional*, defaults to 3):
            Number of decoder layers.
        select_layer_id (`int`, *optional*, defaults to -2):
            Index of the teacher layer to select features from.
        rand_scale (`bool`, *optional*, defaults to `True`):
            Whether to use random scaling during training.
    """

    model_type = "tatok"

    _default_bottleneck = {
        "name": "bottleneck",
        "args": {
            "bottleneck_dim": 1536,
            "norm": "none",
            "regularizer": {
                "name": "simvq",
                "args": {
                    "codebook_loss_weight": 1.0,
                    "codebook_size": 65536,
                    "commitment_loss_weight": 0.25,
                    "entropy_loss_temperature": 0.01,
                    "entropy_loss_weight": 0.0,
                    "l2_normalized": True,
                    "residual_weight": 0.1,
                    "stochastic": True,
                    "stochastic_temperature": 0.03,
                    "top_k": 4,
                    "top_k_prob": 0.5,
                },
            },
        },
    }

    def __init__(
        self,
        bottleneck: Optional[dict] = None,
        bottleneck_token_num: int = 729,
        input_size: int = 384,
        teacher: str = "google/siglip2-so400m-patch14-384",
        input_type: str = "indices",
        pool_scale: int = 1,
        decoder_depth: int = 3,
        select_layer_id: int = -2,
        rand_scale: bool = True,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.bottleneck = bottleneck if bottleneck is not None else copy.deepcopy(self._default_bottleneck)
        self.bottleneck_token_num = bottleneck_token_num
        self.input_size = input_size
        self.teacher = teacher
        self.input_type = input_type
        self.pool_scale = pool_scale
        self.decoder_depth = decoder_depth
        self.select_layer_id = select_layer_id
        self.rand_scale = rand_scale


AutoConfig.register("tatok", TATokConfig)

__all__ = ["TATokConfig"]