# Apache 2.0 - port of togethercomputer's StripedHyenaConfig.
"""Configuration for Evo1 (StripedHyena 7B family)."""

from __future__ import annotations

import json

from transformers import PretrainedConfig


class Evo1Config(PretrainedConfig):
    """Configuration for the Evo1 family.

    Defaults match the evo-1-8k-base / evo-1-131k-base / evo-1.5-8k-base
    checkpoints (32 layers, 4096 hidden, 32 heads, attn at idx [8, 16, 24]
    and Hyena everywhere else, byte-level vocab_size=512). The 131k variant
    overrides ``use_interpolated_rotary_pos_emb`` and ``rotary_emb_scaling_factor``
    plus a longer ``max_seqlen``.
    """

    model_type = "evo1"

    def __init__(
        self,
        # Architecture
        vocab_size: int = 512,
        hidden_size: int = 4096,
        num_filters: int = 4096,
        inner_mlp_size: int = 10928,
        attn_layer_idxs=None,
        hyena_layer_idxs=None,
        num_layers: int = 32,
        num_attention_heads: int = 32,
        proj_groups: int = 1,
        hyena_filter_groups: int = 1,
        short_filter_length: int = 3,
        short_filter_bias: bool = True,
        state_size: int = 8,
        column_split: bool = False,
        column_split_hyena: bool = True,
        split_k0: bool = True,
        smeared_gqa: bool = False,
        # Norms
        eps: float = 1e-6,
        final_norm: bool = True,
        # Linear biases
        mha_out_proj_bias: bool = True,
        qkv_proj_bias: bool = True,
        # Embeddings
        tie_embeddings: bool = True,
        make_vocab_size_divisible_by: int = 8,
        # Activations
        mlp_activation: str = "gelu",
        # Sequence length / RoPE
        max_seqlen: int = 8192,
        rotary_emb_base: float = 10000,
        use_interpolated_rotary_pos_emb: bool = False,
        rotary_emb_scaling_factor: float = 1.0,
        # Inference engine
        prefill_style: str = "fft",
        inference_mode: bool = False,
        # Backend toggles
        use_cache: bool = True,
        use_flash_attention_2: bool = True,
        use_flash_rmsnorm: bool = False,
        use_flash_depthwise: bool = False,
        use_flashfft: bool = False,
        use_flash_attn: bool = False,
        # Misc
        log_intermediate_values: bool = False,
        model_parallel_size: int = 1,
        pipe_parallel_size: int = 1,
        **kwargs,
    ):
        if attn_layer_idxs is None:
            attn_layer_idxs = [8, 16, 24]
        if hyena_layer_idxs is None:
            hyena_layer_idxs = [i for i in range(num_layers) if i not in attn_layer_idxs]

        # Architecture
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_filters = num_filters
        self.inner_mlp_size = inner_mlp_size
        self.attn_layer_idxs = attn_layer_idxs
        self.hyena_layer_idxs = hyena_layer_idxs
        self.num_layers = num_layers
        self.num_attention_heads = num_attention_heads
        self.proj_groups = proj_groups
        self.hyena_filter_groups = hyena_filter_groups
        self.short_filter_length = short_filter_length
        self.short_filter_bias = short_filter_bias
        self.state_size = state_size
        self.column_split = column_split
        self.column_split_hyena = column_split_hyena
        self.split_k0 = split_k0
        self.smeared_gqa = smeared_gqa
        # Norms
        self.eps = eps
        self.final_norm = final_norm
        # Biases
        self.mha_out_proj_bias = mha_out_proj_bias
        self.qkv_proj_bias = qkv_proj_bias
        # Embeddings
        self.tie_embeddings = tie_embeddings
        self.make_vocab_size_divisible_by = make_vocab_size_divisible_by
        # Activations
        self.mlp_activation = mlp_activation
        # Length / RoPE
        self.max_seqlen = max_seqlen
        self.rotary_emb_base = rotary_emb_base
        self.use_interpolated_rotary_pos_emb = use_interpolated_rotary_pos_emb
        self.rotary_emb_scaling_factor = rotary_emb_scaling_factor
        # Engine
        self.prefill_style = prefill_style
        self.inference_mode = inference_mode
        # Backend toggles
        self.use_cache = use_cache
        self.use_flash_attention_2 = use_flash_attention_2
        self.use_flash_rmsnorm = use_flash_rmsnorm
        self.use_flash_depthwise = use_flash_depthwise
        self.use_flashfft = use_flashfft
        self.use_flash_attn = use_flash_attn
        # Misc
        self.log_intermediate_values = log_intermediate_values
        self.model_parallel_size = model_parallel_size
        self.pipe_parallel_size = pipe_parallel_size
        super().__init__(**kwargs)

    # ------------------------------------------------------------------
    # Backwards-compatible attribute access.
    #
    # The internal blocks (RMSNorm, ParallelGatedMLP, ...) call
    # ``config.get(key, default)`` because they were originally written
    # against a `dotdict`. PretrainedConfig has a different `.get`, so we
    # provide a dict-like one that delegates to attribute access.
    # ------------------------------------------------------------------
    @property
    def num_hidden_layers(self) -> int:
        # HF generation utilities (DynamicCache, etc.) expect this name; we
        # keep ``num_layers`` as the source of truth to match the upstream
        # StripedHyena config.
        return self.num_layers

    def get(self, key, default=None):
        # Dict-style access used by internal blocks (RMSNorm, MHA, ...).
        return getattr(self, key, default)

    @classmethod
    def from_original_config(cls, config_path: str, **kwargs) -> "Evo1Config":
        with open(config_path, "r") as f:
            config = json.load(f)
        return cls(**config, **kwargs)