"""Text-only schema25 inference configuration."""
from __future__ import annotations
from collections.abc import Sequence
from typing import Any
from transformers import PreTrainedConfig
from transformers.models.diffusion_gemma import DiffusionGemmaTextConfig

DENOISE_TEMPERATURE = 0.8

class ModilifyMk2TextConfig(DiffusionGemmaTextConfig):
    model_type = "modilify_mk2_text"
    vocab_size: int = 262_144
    hidden_size: int = 2816
    intermediate_size: int = 2112
    num_hidden_layers: int = 30
    num_attention_heads: int = 16
    num_key_value_heads: int = 8
    head_dim: int = 256
    max_position_embeddings: int = 262_144
    sliding_window: int = 1024
    use_bidirectional_attention: str | None = None
    num_global_key_value_heads: int | None = 2
    global_head_dim: int = 512
    num_experts: int | None = 128
    top_k_experts: int | None = 8
    moe_intermediate_size: int | None = 704

class ModilifyMk2Config(PreTrainedConfig):
    model_type = "modilify_mk2"
    sub_configs = {"text_config": ModilifyMk2TextConfig}

    def __init__(
        self, text_config: ModilifyMk2TextConfig | dict[str, Any] | None = None, *,
        canvas_length: int = 256,
        initializer_range: float = 0.02,
        tie_word_embeddings: bool = True,
        state_schema_version: int = 25,
        memory_architecture: str = "compact_gdn2_v2",
        latent_dim: int = 2816,
        latent_ffn_dim: int = 7168,
        latent_num_layers: int = 4,
        latent_num_heads: int = 16,
        latent_local_attention_window: int = 128,
        latent_tape_probes: int = 4,
        latent_history_kv_rank: int = 1024,
        latent_working_last_block_global: bool = True,
        working_memory_bus: bool = True,
        persistent_memory_bus: bool = True,
        commit_sequence_dim: int = 1024,
        kv_cache_bucket_size: int = 128,
        vocab_chunk_size: int = 32768,
        turn_end_token_id: int = 106,
        terminal_token_ids: Sequence[int] = (106, 50),
        eos_token_id: int = 1,
        commit_failure_budget: float = 0.2,
        commit_top_k: int | None = 40,
        commit_min_p: float | None = 0.05,
        commit_target_confidence: float | None = 0.5,
        commit_entropy_weight: float = 1.0,
        commit_confidence_power: float = 1.0,
        **kwargs: Any,
    ) -> None:
        if state_schema_version != 25 or memory_architecture != "compact_gdn2_v2":
            raise ValueError("This release requires schema25 compact_gdn2_v2 weights.")
        if text_config is None:
            text_config = ModilifyMk2TextConfig()
        elif isinstance(text_config, dict):
            payload = dict(text_config)
            payload.pop("model_type", None)
            text_config = ModilifyMk2TextConfig(**payload)
        text_config.use_bidirectional_attention = None
        self.text_config = text_config
        self.canvas_length = canvas_length
        self.initializer_range = initializer_range
        self.state_schema_version = state_schema_version
        self.memory_architecture = memory_architecture
        self.latent_dim = latent_dim
        self.latent_ffn_dim = latent_ffn_dim
        self.latent_num_layers = latent_num_layers
        self.latent_num_heads = latent_num_heads
        self.latent_local_attention_window = latent_local_attention_window
        self.latent_tape_probes = latent_tape_probes
        self.latent_history_kv_rank = latent_history_kv_rank
        self.latent_working_last_block_global = latent_working_last_block_global
        self.working_memory_bus = working_memory_bus
        self.persistent_memory_bus = persistent_memory_bus
        self.commit_sequence_dim = commit_sequence_dim
        self.kv_cache_bucket_size = kv_cache_bucket_size
        self.vocab_chunk_size = vocab_chunk_size
        self.turn_end_token_id = turn_end_token_id
        self.terminal_token_ids = terminal_token_ids
        self.commit_failure_budget = commit_failure_budget
        self.commit_top_k = commit_top_k
        self.commit_min_p = commit_min_p
        self.commit_target_confidence = commit_target_confidence
        self.commit_entropy_weight = commit_entropy_weight
        self.commit_confidence_power = commit_confidence_power
        super().__init__(tie_word_embeddings=tie_word_embeddings,
                         eos_token_id=eos_token_id, **kwargs)