from __future__ import annotations

from dataclasses import asdict, dataclass

MODEL_NAME = "Sol Milkshake"
ORGANIZATION = "Sol Labs"
DEPLOYED_PARAMS = 2_990_000
VOCAB_SIZE = 2_048
D_MODEL = 192
MAX_CONTEXT = 2_048
N_PHYSICAL_BLOCKS = 5
N_Q_HEADS = 6
N_KV_HEADS = 2
HEAD_DIM = 32
FFN_HIDDEN = 512
TN_RANK = 26
TIED_EMBEDDING_HEAD = True
XSA_PARAMETER_COUNT = 0
TOKENS_PER_PARAMETER = 544
GLOBAL_TOKENS = 1_626_560_000
GLOBAL_BATCH = 32_768
COMPLETE_UPDATES = 49_638
FINAL_PARTIAL_TOKENS = 22_016

assert DEPLOYED_PARAMS * TOKENS_PER_PARAMETER == GLOBAL_TOKENS
assert COMPLETE_UPDATES * GLOBAL_BATCH + FINAL_PARTIAL_TOKENS == GLOBAL_TOKENS


@dataclass(frozen=True)
class SolConfig:
    model_name: str = MODEL_NAME
    vocab_size: int = VOCAB_SIZE
    d_model: int = D_MODEL
    max_context: int = MAX_CONTEXT
    n_blocks: int = N_PHYSICAL_BLOCKS
    n_q_heads: int = N_Q_HEADS
    n_kv_heads: int = N_KV_HEADS
    head_dim: int = HEAD_DIM
    ffn_hidden: int = FFN_HIDDEN
    tn_rank: int = TN_RANK
    rope_theta: float = 20_000.0
    memory_width: int = 64
    memory_slots: int = 32
    chunk_size: int = 32
    passes: int = 3
    dropout: float = 0.0

    def as_dict(self) -> dict:
        return asdict(self)