from dataclasses import dataclass, field
from typing import Optional


@dataclass
class SharonaConfig:
    vocab_size: int = 32_768
    hidden_size: int = 3_072
    num_self_decoder_layers: int = 14
    num_cross_decoder_layers: int = 14
    num_query_heads: int = 24
    num_kv_heads: int = 8
    head_dim: int = 128
    ffn_intermediate_size: int = 8_192
    max_context1_length: int = 152_000
    max_context2_length: int = 228_000
    sliding_window_size: int = 4_096
    rope_theta: float = 500_000.0
    rope_scaling_factor: Optional[float] = None
    rope_scaling_type: Optional[str] = None
    rope_low_freq_factor: float = 1.0
    rope_high_freq_factor: float = 4.0
    rope_original_max_position_embeddings: int = 8_192
    rms_norm_eps: float = 1e-6
    initializer_range: float = 0.02
    tie_word_embeddings: bool = True
    use_gradient_checkpointing: bool = True
    kv_cache_dtype: str = "bfloat16"
    num_persistent_contexts: int = 1
    m1_length: int = 512
    m3_length: int = 1536
    attention_bias: bool = False
    pad_token_id: int = 32_764

    @property
    def total_layers(self) -> int:
        return self.num_self_decoder_layers + self.num_cross_decoder_layers

    @property
    def gqa_ratio(self) -> int:
        return self.num_query_heads // self.num_kv_heads

    @classmethod
    def sharona_3b(cls) -> "SharonaConfig":
        return cls()

    @classmethod
    def sharona_3b_triple(cls) -> "SharonaConfig":
        # Fixed architecture constants for the 3B² triple-context build.
        return cls(
            vocab_size=128_256,
            rms_norm_eps=1e-5,
            rope_scaling_type="scaled",
            rope_scaling_factor=32.0,
            rope_low_freq_factor=1.0,
            rope_high_freq_factor=4.0,
            rope_original_max_position_embeddings=8_192,
            pad_token_id=128_004,
            num_persistent_contexts=2,
            m1_length=512,
            m3_length=1536,
            max_context1_length=2048,
        )

    @classmethod
    def sharona_8b_triple(cls) -> "SharonaConfig":
        return cls(
            vocab_size=128_256,
            hidden_size=4_096,
            num_self_decoder_layers=16,
            num_cross_decoder_layers=16,
            num_query_heads=32,
            num_kv_heads=8,
            head_dim=128,
            ffn_intermediate_size=14_336,
            max_context1_length=131_072,
            max_context2_length=228_000,
            sliding_window_size=4_096,
            rope_theta=500_000.0,
            rope_scaling_type="scaled",
            rope_scaling_factor=32.0,
            rope_low_freq_factor=1.0,
            rope_high_freq_factor=4.0,
            rope_original_max_position_embeddings=8_192,
            rms_norm_eps=1e-5,
            initializer_range=0.02,
            tie_word_embeddings=False,
            num_persistent_contexts=2,
            m1_length=512,
            m3_length=2048,
            pad_token_id=128_004,
        )

    @classmethod
    def sharona_36b_triple(cls) -> "SharonaConfig":
        return cls(
            vocab_size=155_136,
            hidden_size=5_120,
            num_self_decoder_layers=32,
            num_cross_decoder_layers=32,
            num_query_heads=80,
            num_kv_heads=8,
            head_dim=128,
            ffn_intermediate_size=27_648,
            max_context1_length=524_288,
            max_context2_length=524_288,
            sliding_window_size=4_096,
            rope_theta=10_000_000.0,
            rope_scaling_type=None,
            rope_scaling_factor=None,
            rms_norm_eps=1e-6,
            initializer_range=0.02,
            tie_word_embeddings=False,
            attention_bias=True,
            num_persistent_contexts=2,
            m1_length=512,
            m3_length=2048,
            pad_token_id=1,
        )

    @classmethod
    def sharona_8b(cls) -> "SharonaConfig":
        return cls(
            vocab_size=32_768,
            hidden_size=4_096,
            num_self_decoder_layers=16,
            num_cross_decoder_layers=16,
            num_query_heads=32,
            num_kv_heads=8,
            head_dim=128,
            ffn_intermediate_size=14_336,
            max_context1_length=131_072,
            max_context2_length=228_000,
            sliding_window_size=4_096,
            rope_theta=500_000.0,
            rms_norm_eps=1e-5,
            initializer_range=0.02,
            tie_word_embeddings=False,
        )

    @classmethod
    def sharona_7b(cls) -> "SharonaConfig":
        """B² from Qwen2.5-7B-Instruct: 14+14 layers, triple context."""
        return cls(
            vocab_size=152_064,
            hidden_size=3_584,
            num_self_decoder_layers=14,
            num_cross_decoder_layers=14,
            num_query_heads=28,
            num_kv_heads=4,
            head_dim=128,
            ffn_intermediate_size=18_944,
            max_context1_length=32_768,
            max_context2_length=32_768,
            sliding_window_size=4_096,
            rope_theta=1_000_000.0,
            rms_norm_eps=1e-6,
            attention_bias=True,
            initializer_range=0.02,
            tie_word_embeddings=False,
            use_gradient_checkpointing=True,
            num_persistent_contexts=2,
            m1_length=2_048,
            m3_length=16_384,
            pad_token_id=151_643,
        )