"""
DeepSeek-V4 Configuration
=========================
HuggingFace-compatible configuration for the DeepSeek-V4 architecture.

Key innovations over DeepSeek-V3:
- Compressed Sparse Attention (CSA) with sliding window + learned KV compression
- Manifold-Constrained Hyper-Connections (mHC) replacing residual connections  
- Hash routing for first N layers + sqrtsoftplus scoring
- Grouped low-rank output projection
- Attention sink (learnable per-head bias)
"""

from transformers import PretrainedConfig


class DeepSeekV4Config(PretrainedConfig):
    """Configuration class for the DeepSeek-V4 model.
    
    This config captures ALL architectural parameters from the DeepSeek-V4 technical report
    and official inference code, enabling faithful reproduction.
    """
    
    model_type = "deepseek_v4"
    
    def __init__(
        self,
        # Core dimensions
        vocab_size: int = 129280,
        hidden_size: int = 7168,
        num_hidden_layers: int = 61,
        num_attention_heads: int = 128,
        
        # MLA (Multi-head Latent Attention) parameters
        q_lora_rank: int = 1536,
        head_dim: int = 512,            # unified KV head dim (V4 unifies nope+rope+v)
        qk_rope_head_dim: int = 64,     # decoupled RoPE dimension
        num_key_value_heads: int = 1,    # effectively MQA at compression level
        
        # Grouped output projection
        o_groups: int = 16,              # number of output groups
        o_lora_rank: int = 1024,         # per-group output compression rank
        
        # Compressed Sparse Attention
        window_size: int = 128,          # sliding window size
        compress_ratios: tuple = None,   # per-layer compression ratios (0=no compression)
        compress_rope_theta: float = 160000.0,  # RoPE theta for compressed KV
        
        # Indexer (learned sparse attention selection)
        index_n_heads: int = 64,
        index_head_dim: int = 128,
        index_topk: int = 1024,
        
        # Attention sink
        attention_bias: bool = False,    # standard bias in linear layers
        attention_dropout: float = 0.0,
        
        # MoE parameters
        moe_intermediate_size: int = 3072,
        n_routed_experts: int = 384,
        n_shared_experts: int = 1,
        num_experts_per_tok: int = 6,
        num_hash_layers: int = 3,        # first N layers use hash routing
        scoring_func: str = "sqrtsoftplus",  # NEW in V4: sqrt(softplus(x))
        topk_method: str = "noaux_tc",
        norm_topk_prob: bool = True,
        routed_scaling_factor: float = 2.5,
        swiglu_limit: float = 10.0,      # NEW in V4: clamp SwiGLU activations
        
        # Expert quantization
        expert_dtype: str = "fp4",       # NEW in V4: FP4 expert weights
        
        # Hyper-Connections (mHC)
        hc_mult: int = 4,               # number of hidden state copies
        hc_sinkhorn_iters: int = 20,     # Sinkhorn normalization iterations
        hc_eps: float = 1e-6,
        
        # Multi-Token Prediction
        num_nextn_predict_layers: int = 1,
        
        # Normalization
        rms_norm_eps: float = 1e-6,
        hidden_act: str = "silu",
        
        # RoPE parameters
        rope_theta: float = 10000.0,
        rope_scaling: dict = None,
        max_position_embeddings: int = 1048576,  # 1M context
        
        # Quantization config (for inference)
        quantization_config: dict = None,
        
        # Standard HF config
        initializer_range: float = 0.02,
        use_cache: bool = True,
        tie_word_embeddings: bool = False,
        bos_token_id: int = 0,
        eos_token_id: int = 1,
        torch_dtype: str = "bfloat16",
        
        **kwargs,
    ):
        super().__init__(
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )
        
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        
        # MLA
        self.q_lora_rank = q_lora_rank
        self.head_dim = head_dim
        self.qk_rope_head_dim = qk_rope_head_dim
        self.num_key_value_heads = num_key_value_heads
        
        # Grouped output
        self.o_groups = o_groups
        self.o_lora_rank = o_lora_rank
        
        # Compressed Sparse Attention
        self.window_size = window_size
        if compress_ratios is not None:
            self.compress_ratios = list(compress_ratios)
        else:
            # Default: V4-Pro pattern — alternating 128 and 4 compression ratios
            # First 2 layers: 128 (HCA only), then alternating 4,128 pairs, last layer: 0
            self.compress_ratios = [128, 128] + [4, 128] * 29 + [4, 0]
        self.compress_rope_theta = compress_rope_theta
        
        # Indexer
        self.index_n_heads = index_n_heads
        self.index_head_dim = index_head_dim
        self.index_topk = index_topk
        
        # Attention
        self.attention_bias = attention_bias
        self.attention_dropout = attention_dropout
        
        # MoE
        self.moe_intermediate_size = moe_intermediate_size
        self.n_routed_experts = n_routed_experts
        self.n_shared_experts = n_shared_experts
        self.num_experts_per_tok = num_experts_per_tok
        self.num_hash_layers = num_hash_layers
        self.scoring_func = scoring_func
        self.topk_method = topk_method
        self.norm_topk_prob = norm_topk_prob
        self.routed_scaling_factor = routed_scaling_factor
        self.swiglu_limit = swiglu_limit
        self.expert_dtype = expert_dtype
        
        # Hyper-Connections
        self.hc_mult = hc_mult
        self.hc_sinkhorn_iters = hc_sinkhorn_iters
        self.hc_eps = hc_eps
        
        # MTP
        self.num_nextn_predict_layers = num_nextn_predict_layers
        
        # Normalization
        self.rms_norm_eps = rms_norm_eps
        self.hidden_act = hidden_act
        
        # RoPE
        self.rope_theta = rope_theta
        self.rope_scaling = rope_scaling or {
            "type": "yarn",
            "factor": 16,
            "original_max_position_embeddings": 65536,
            "beta_fast": 32,
            "beta_slow": 1,
        }
        self.max_position_embeddings = max_position_embeddings
        
        # Quantization
        self.quantization_config = quantization_config
        
        # Standard
        self.initializer_range = initializer_range
        self.use_cache = use_cache
        self.torch_dtype = torch_dtype