# !!!!!!!!!!!!!!! RiXIS 1 [PREVIEW] !!!!!!!!!!!!!!!
# Authorised public RiXIS 1 model weights release ("NeuraNET Zero").
# Source files are a reference implementation for loading and
# inference. proprietary development infrastructure and implementation
# details are omitted.
#
# Copyright (c) 2026 Ruben Roy. All rights reserved.
#
# Licensed under the Creative Commons Attribution-NonCommercial-
# NoDerivatives 4.0 International License (CC BY-NC-ND 4.0);
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://creativecommons.org/licenses/by-nc-nd/4.0/
#
# Unless required by applicable law or agreed to in writing, this work
# 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.

from __future__ import annotations

from transformers import PreTrainedConfig


class RiXIS1Config(PreTrainedConfig):

    # /\/\ RiXIS 1 decoder-only language model /\/\

    model_type = "rixis1"
    keys_to_ignore_at_inference = ["past_key_values"]

    def __init__(
        self,
        vocab_size: int = 32_001,
        hidden_size: int = 4_096,
        intermediate_size: int = 14_336,
        num_hidden_layers: int = 80,
        num_attention_heads: int = 32,
        num_key_value_heads: int = 8,
        head_dim: int | None = 128,
        hidden_act: str = "silu",
        max_position_embeddings: int = 32_768,
        initializer_range: float = 0.02,
        rms_norm_eps: float = 1e-5,
        use_cache: bool = True,
        pad_token_id: int | None = 32_000,
        bos_token_id: int | None = 1,
        eos_token_id: int | list[int] | None = 2,
        tie_word_embeddings: bool = False,
        rope_parameters: dict | None = None,
        rope_theta: float | None = None,
        sliding_window: int | None = None,
        attention_dropout: float = 0.0,
        **kwargs,
    ):
        if hidden_size % num_attention_heads != 0:
            raise ValueError(
                "hidden_size must be divisible by num_attention_heads."
            )
        if num_attention_heads % num_key_value_heads != 0:
            raise ValueError(
                "num_attention_heads must be divisible by num_key_value_heads."
            )

        inferred_head_dim = hidden_size // num_attention_heads
        if head_dim is None:
            head_dim = inferred_head_dim
        if head_dim != inferred_head_dim:
            raise ValueError(
                f"head_dim={head_dim} is incompatible with hidden_size="
                f"{hidden_size} and num_attention_heads={num_attention_heads}."
            )

        if rope_parameters is None:
            rope_parameters = {
                "rope_type": "default",
                "rope_theta": float(rope_theta or 10_000.0),
            }
        else:
            rope_parameters = dict(rope_parameters)
            rope_parameters.setdefault("rope_type", "default")
            if "rope_theta" not in rope_parameters:
                rope_parameters["rope_theta"] = float(
                    rope_theta or 10_000.0
                )

        if rope_parameters["rope_type"] != "default":
            raise ValueError(
                "PUBLIC RiXIS (x1x) arch currently supports "
                "the default rotary-position formulation only."
            )

        layer_types = kwargs.pop(
            "layer_types",
            ["full_attention"] * num_hidden_layers,
        )
        if len(layer_types) != num_hidden_layers:
            raise ValueError(
                "layer_types must contain exactly num_hidden_layers entries."
            )

        kwargs.setdefault("is_decoder", True)
        kwargs.setdefault("is_encoder_decoder", False)

        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.head_dim = head_dim
        self.hidden_act = hidden_act
        self.max_position_embeddings = max_position_embeddings
        self.initializer_range = initializer_range
        self.rms_norm_eps = rms_norm_eps
        self.use_cache = use_cache
        self.rope_parameters = rope_parameters
        self.sliding_window = sliding_window
        self.attention_dropout = attention_dropout
        self.layer_types = layer_types

        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )


__all__ = ["RiXIS1Config"]