#!/usr/bin/env python3
"""
Dhara-AR: Configuration for Dhara autoregressive language model.

LLaMA3-style architecture with Canon Layer positions (ABCD) from
"Physics of Language Models: Part 4.1" by Zeyuan Allen-Zhu.
"""

from transformers import PretrainedConfig


class DharaARConfig(PretrainedConfig):
    """Configuration for Dhara-AR model."""

    model_type = "dhara_ar"

    def __init__(
        self,
        # Core architecture - ~250M params
        vocab_size: int = 49152,
        hidden_size: int = 768,
        intermediate_size: int = 2176,    # Tuned for 250M total params
        num_hidden_layers: int = 32,
        num_attention_heads: int = 12,
        num_key_value_heads: int = 4,
        max_position_embeddings: int = 8192,

        # Model specifics
        hidden_act: str = "silu",
        rms_norm_eps: float = 1e-6,
        rope_theta: float = 100000.0,
        initializer_range: float = 0.02,
        tie_word_embeddings: bool = True,
        attention_bias: bool = False,
        attention_dropout: float = 0.0,
        mlp_bias: bool = False,

        # Enhancements
        use_qk_norm: bool = True,
        use_logit_softcap: bool = True,
        logit_softcap: float = 30.0,

        # RoPE scaling (for inference-time context extension)
        rope_scaling: dict = None,  # {"type": "yarn", "factor": 2.0} for 2x extension

        # Canon layer parameters - ALL 4 positions
        canon_set: str = "ABCD",
        canon_kernel: int = 4,
        canon_residual: bool = True,
        canon_activation: bool = False,
        canon_bias: bool = False,

        **kwargs
    ):
        super().__init__(**kwargs)
        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.max_position_embeddings = max_position_embeddings
        self.hidden_act = hidden_act
        self.rms_norm_eps = rms_norm_eps
        self.rope_theta = rope_theta
        self.initializer_range = initializer_range
        self.tie_word_embeddings = tie_word_embeddings
        self.attention_bias = attention_bias
        self.attention_dropout = attention_dropout
        self.mlp_bias = mlp_bias

        # Enhancements
        self.use_qk_norm = use_qk_norm
        self.use_logit_softcap = use_logit_softcap
        self.logit_softcap = logit_softcap

        # RoPE scaling
        self.rope_scaling = rope_scaling

        # Canon config
        self.canon_set = canon_set
        self.canon_kernel = canon_kernel
        self.canon_residual = canon_residual
        self.canon_activation = canon_activation
        self.canon_bias = canon_bias