"""SeqLens model configuration — v2."""

from dataclasses import dataclass


@dataclass
class SeqLensConfig:
    """Configuration for SeqLens genomic language model.

    Micro v2: 8 layers, dim 256, attention every 4th layer, ~10M params, 16K context.
    """

    # Vocabulary: A=0, T=1, G=2, C=3, N=4, [CLS]=5, [SEP]=6, [PAD]=7, [MASK]=8
    vocab_size: int = 9
    pad_token_id: int = 7
    mask_token_id: int = 8
    cls_token_id: int = 5
    sep_token_id: int = 6

    # Model dimensions
    d_model: int = 256
    n_layers: int = 8           # v2: doubled from 4

    # Mamba2 SSM
    ssm_d_state: int = 64
    ssm_d_conv: int = 4
    ssm_expand: int = 2
    ssm_headdim: int = 64

    # Sliding-window attention
    n_attn_heads: int = 4
    attn_window: int = 512
    attn_layer_interval: int = 4  # Attention at layers 3, 7 (0-indexed)

    # Feed-forward
    ffn_expand: int = 4

    # Sequence
    max_seq_len: int = 16_384

    # Regularization
    dropout: float = 0.0

    # Training
    mask_rate: float = 0.15
    mask_token_prob: float = 0.80
    mask_random_prob: float = 0.10

    @property
    def d_inner(self):
        return self.d_model * self.ssm_expand

    @property
    def n_ssm_heads(self):
        return self.d_inner // self.ssm_headdim

    @property
    def attn_head_dim(self):
        return self.d_model // self.n_attn_heads

    def has_attention(self, layer_idx: int) -> bool:
        if self.attn_layer_interval <= 0:
            return False
        return (layer_idx + 1) % self.attn_layer_interval == 0


MICRO_CONFIG = SeqLensConfig(
    d_model=256,
    n_layers=8,
    max_seq_len=16_384,
)

BASE_CONFIG = SeqLensConfig(
    d_model=512,
    n_layers=12,
    n_attn_heads=8,
    max_seq_len=32_768,
)

LARGE_CONFIG = SeqLensConfig(
    d_model=768,
    n_layers=24,
    n_attn_heads=12,
    max_seq_len=65_536,
)