from __future__ import annotations

from typing import Any

from transformers import LlamaConfig, PretrainedConfig


class NcpSmolConfig(PretrainedConfig):
    model_type = "ncp_smol"

    def __init__(
        self,
        backbone_config: dict[str, Any] | None = None,
        base_model_name_or_path: str = "HuggingFaceTB/SmolLM2-360M",
        base_model_revision: str | None = None,
        chunk_size: int = 4,
        segments: int | None = None,
        codebook_size: int = 64,
        concept_layers: int = 2,
        insert_layer: int = 1,
        ncp_target: str = "continuous",
        ncp_weight: float = 1.0,
        vq_weight: float = 1.0,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        backbone_config = backbone_config or LlamaConfig().to_dict()
        hidden_size = int(backbone_config["hidden_size"])
        heads = int(backbone_config["num_attention_heads"])
        segments = heads if segments is None else segments

        if chunk_size < 2:
            raise ValueError("chunk_size must be at least 2")
        if hidden_size % segments:
            raise ValueError("hidden_size must be divisible by segments")
        if not 0 < insert_layer < int(backbone_config["num_hidden_layers"]):
            raise ValueError("insert_layer must split the token backbone")
        if concept_layers < 1:
            raise ValueError("concept_layers must be positive")
        if codebook_size < 2:
            raise ValueError("codebook_size must be at least 2")
        if ncp_target not in {"quantized", "continuous"}:
            raise ValueError("ncp_target must be 'quantized' or 'continuous'")

        self.backbone_config = backbone_config
        self.hidden_size = hidden_size
        self.vocab_size = int(backbone_config["vocab_size"])
        self.num_hidden_layers = int(backbone_config["num_hidden_layers"])
        self.num_attention_heads = heads
        self.num_key_value_heads = int(backbone_config["num_key_value_heads"])
        self.use_cache = False
        self.base_model_name_or_path = base_model_name_or_path
        self.base_model_revision = base_model_revision
        self.chunk_size = chunk_size
        self.segments = segments
        self.codebook_size = codebook_size
        self.concept_layers = concept_layers
        self.insert_layer = insert_layer
        self.ncp_target = ncp_target
        self.ncp_weight = ncp_weight
        self.vq_weight = vq_weight
        self.architectures = ["NcpSmolForCausalLM"]
        self.auto_map = {
            "AutoConfig": "configuration.NcpSmolConfig",
            "AutoModelForCausalLM": "modeling.NcpSmolForCausalLM",
        }


NcpSmolConfig.register_for_auto_class()