from transformers import PretrainedConfig


class OrionT2Config(PretrainedConfig):
    """Configuration for Project Prism's Orion T2 / Orion Flagship Mini-T2."""

    model_type = "orion_t2"

    def __init__(
        self,
        vocab_size=32768,
        d_model=768,
        n_layers=12,
        n_heads=12,
        n_kv_heads=4,
        rope_theta=10000.0,
        sequence_length=1024,
        n_procedure_banks=3,
        procedure_bank_span=4,
        n_experts=16,
        top_k=2,
        expert_hidden=1024,
        shared_expert=True,
        vault_key_parts=256,
        vault_slots=65536,
        vault_top_component=12,
        vault_top_k=4,
        vault_read_layers=(3, 7, 11),
        working_state_layers=(3, 7, 11),
        deliberation_start=9,
        deliberation_cycles=1,
        router_aux_coef=0.005,
        router_z_coef=0.0005,
        rms_norm_eps=1e-6,
        use_cache=False,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.rope_theta = rope_theta
        self.sequence_length = sequence_length

        self.n_procedure_banks = n_procedure_banks
        self.procedure_bank_span = procedure_bank_span
        self.n_experts = n_experts
        self.top_k = top_k
        self.expert_hidden = expert_hidden
        self.shared_expert = shared_expert

        self.vault_key_parts = vault_key_parts
        self.vault_slots = vault_slots
        self.vault_top_component = vault_top_component
        self.vault_top_k = vault_top_k
        self.vault_read_layers = list(vault_read_layers)
        self.working_state_layers = list(working_state_layers)

        self.deliberation_start = deliberation_start
        self.deliberation_cycles = deliberation_cycles

        self.router_aux_coef = router_aux_coef
        self.router_z_coef = router_z_coef
        self.rms_norm_eps = rms_norm_eps

        # Standard Transformers aliases.
        self.hidden_size = d_model
        self.num_hidden_layers = n_layers
        self.num_attention_heads = n_heads
        self.num_key_value_heads = n_kv_heads
        self.max_position_embeddings = sequence_length

        # Orion v1 inference intentionally recomputes the full sequence.
        # KV caching is not exposed yet because recurrent deliberation reuses
        # late Cortex stages and Working State, so a naive standard cache would
        # not be equivalent to the training graph.
        self.use_cache = False

        kwargs.setdefault("tie_word_embeddings", False)
        super().__init__(**kwargs)