"""MinSparkConfig: Transformers PretrainedConfig wrapper around MeiosisConfig.

min-spark is the public-facing name of the Meiosis decay-p09 model. This
config carries every MeiosisConfig field (the training config stays a
dataclass in models/2026-07-meiosis/meiosis.py) plus the inference-only
fields effort/use_cache and the standard max_position_embeddings that
transformers and lm-eval read for max length.
"""
from __future__ import annotations

from transformers import PretrainedConfig


class MinSparkConfig(PretrainedConfig):
    model_type = "minspark"

    def __init__(
        self,
        vocab_size: int = 4096,
        dim: int = 288,
        n_heads: int = 6,
        n_kv_heads: int = 2,
        ffn_hidden: int = 768,
        prelude_layers: int = 1,
        coda_layers: int = 1,
        body_blocks: int = 3,
        max_loops: int = 4,
        train_loops: int = 3,
        lora_rank: int = 16,
        rope_base: float = 10000.0,
        max_seq_len: int = 512,
        ddl_beta_init: float = 1.0,
        ddl_k_eps: float = 1e-2,
        ddl_v_sigmoid_scale: float = 4.0,
        doc_mask_eos: int = 2,
        effort: str = "medium",
        use_cache: bool = False,
        max_position_embeddings: int | None = None,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.dim = dim
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.ffn_hidden = ffn_hidden
        self.prelude_layers = prelude_layers
        self.coda_layers = coda_layers
        self.body_blocks = body_blocks
        self.max_loops = max_loops
        self.train_loops = train_loops
        self.lora_rank = lora_rank
        self.rope_base = rope_base
        self.max_seq_len = max_seq_len
        self.ddl_beta_init = ddl_beta_init
        self.ddl_k_eps = ddl_k_eps
        self.ddl_v_sigmoid_scale = ddl_v_sigmoid_scale
        self.doc_mask_eos = doc_mask_eos
        self.effort = effort
        self.use_cache = use_cache
        # Single source for 512: max_position_embeddings derives from
        # max_seq_len so a future context-length change has one place to edit.
        self.max_position_embeddings = (
            max_position_embeddings if max_position_embeddings is not None else max_seq_len
        )
        self.tie_word_embeddings = True
        super().__init__(**kwargs)

    def to_meiosis(self):
        """Rebuild the training MeiosisConfig dataclass from this config."""
        try:
            from .meiosis import MeiosisConfig  # remote-code: vendored sibling
        except ImportError:
            from meiosis import MeiosisConfig  # direct import

        return MeiosisConfig(
            vocab_size=self.vocab_size,
            dim=self.dim,
            n_heads=self.n_heads,
            n_kv_heads=self.n_kv_heads,
            ffn_hidden=self.ffn_hidden,
            prelude_layers=self.prelude_layers,
            coda_layers=self.coda_layers,
            body_blocks=self.body_blocks,
            max_loops=self.max_loops,
            train_loops=self.train_loops,
            lora_rank=self.lora_rank,
            rope_base=self.rope_base,
            max_seq_len=self.max_seq_len,
            ddl_beta_init=self.ddl_beta_init,
            ddl_k_eps=self.ddl_k_eps,
            ddl_v_sigmoid_scale=self.ddl_v_sigmoid_scale,
            doc_mask_eos=self.doc_mask_eos,
        )

    @classmethod
    def from_meiosis(cls, cfg):
        return cls(
            vocab_size=cfg.vocab_size, dim=cfg.dim, n_heads=cfg.n_heads,
            n_kv_heads=cfg.n_kv_heads, ffn_hidden=cfg.ffn_hidden,
            prelude_layers=cfg.prelude_layers, coda_layers=cfg.coda_layers,
            body_blocks=cfg.body_blocks, max_loops=cfg.max_loops,
            train_loops=cfg.train_loops, lora_rank=cfg.lora_rank,
            rope_base=cfg.rope_base, max_seq_len=cfg.max_seq_len,
            ddl_beta_init=cfg.ddl_beta_init, ddl_k_eps=cfg.ddl_k_eps,
            ddl_v_sigmoid_scale=cfg.ddl_v_sigmoid_scale, doc_mask_eos=cfg.doc_mask_eos,
        )