"""
Qwen3 Model Configuration.
"""

from dataclasses import dataclass, field
from typing import List, Optional


@dataclass
class Qwen3Config:
    """Configuration class for Qwen3 model."""
    vocab_size: int = 151936
    hidden_size: int = 1024
    intermediate_size: int = 3072
    num_hidden_layers: int = 28
    num_attention_heads: int = 16
    num_key_value_heads: int = 8
    head_dim: int = 128
    hidden_act: str = "silu"
    max_position_embeddings: int = 32768
    initializer_range: float = 0.02
    rms_norm_eps: float = 1e-6
    use_cache: bool = True
    tie_word_embeddings: bool = True
    rope_theta: float = 1000000.0
    rope_scaling: Optional[dict] = None
    # ``rope`` (default) preserves legacy behavior; ``alibi`` disables RoPE
    # and applies Press et al. head-specific causal distance biases.
    positional: str = "rope"
    # QK normalisation is part of the published Qwen3-style attention path.
    # Keep it on by default so checkpoints/configs written before this flag are
    # bit-identical; ``off`` removes both per-head RMSNorm modules.
    qk_norm: str = "on"
    # Keep Q/K RMSNorm gains and their AdamW state in FP32 during training.
    qk_gain_fp32: bool = False
    # Training-only RoPE position scales. When non-empty, every attention layer
    # independently samples one scale per forward pass. Evaluation always uses
    # the standard scale 1.0. This provides a dense positional-augmentation
    # control without deleting attention links.
    rope_train_position_scale_set: Optional[List[float]] = None
    attention_bias: bool = False
    attention_dropout: float = 0.0
    bos_token_id: int = 151643
    eos_token_id: int = 151643
    checkpoint_activations: bool = False
    attention_type: str = "dense"
    sparse_block_size: int = 128
    sparse_num_random_blocks: int = 1
    sparse_mask_pool_size: int = 1000
    sparse_max_seq_length: int = 1024
    # When set, only the listed layer indices use sparse attention; all others use dense.
    # When None (default), attention_type applies to all layers uniformly.
    sparse_layer_indices: Optional[List[int]] = None
    # Directory to cache pre-built block masks (one file per layer). If None, masks are rebuilt every run.
    sparse_mask_cache_dir: Optional[str] = None
    # Seed offset for the frozen (pool_size=1) fixed mask. 0 = legacy (layer-index only).
    # Distinct values yield distinct fixed masks (fixed-mask seeds).
    sparse_mask_seed: int = 0
    # Mixed-sparsity training: if set to a non-empty list, each mask-pool entry draws its
    # OWN num_random_blocks uniformly from this set, so the per-step mask (picked at random
    # from the pool) also resamples the sparsity LEVEL. None/empty = legacy fixed-nrb pool.
    # This is the "Mixed" training arm of the paper; without it the arm silently degenerates
    # to a fixed-nrb model.
    sparse_mixed_nrb_set: Optional[List[int]] = None
    # --- deployment-time baselines (learned MoBA + ASEntmax) ---
    # moba_top_k is the per-query block budget used by the top-k deployment ladder.
    moba_block_size: int = 128
    moba_top_k: int = 32
    asentmax_alpha: float = 1.5
    asentmax_learn_temp: bool = True
    asentmax_init_temp: float = 1.0

    def is_layer_sparse(self, layer_idx: int) -> bool:
        """Return True if this layer should use sparse stochastic attention."""
        if self.attention_type != "sparse_stochastic":
            return False
        if self.sparse_layer_indices is None:
            return True  # all layers sparse (existing behaviour)
        return layer_idx in self.sparse_layer_indices

    @classmethod
    def from_dict(cls, config_dict: dict) -> "Qwen3Config":
        return cls(**{k: v for k, v in config_dict.items() if k in cls.__dataclass_fields__})
    
    @classmethod
    def from_json(cls, path: str) -> "Qwen3Config":
        import json
        with open(path, "r") as f:
            config_dict = json.load(f)
        return cls.from_dict(config_dict)