# Copyright 2026 Modilify
# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
"""Configuration classes for Modilify Mk1."""

from __future__ import annotations

import math
from typing import Any

from transformers.models.diffusion_gemma import (
    DiffusionGemmaConfig,
    DiffusionGemmaTextConfig,
)


class ModilifyMk1TextConfig(DiffusionGemmaTextConfig):
    """Text configuration for the Modilify Mk1 decoder.

    This class preserves the standard DiffusionGemma text schema while giving
    the exported model an independent, stable model type.
    """

    model_type = "modilify_mk1_text"


class ModilifyMk1Config(DiffusionGemmaConfig):
    """Serializable multimodal inference configuration for Modilify Mk1.

    Args:
        text_config: DiffusionGemma text configuration or its serialized form.
        vision_config: Gemma 4 vision configuration or its serialized form.
        denoise_temperature: Sampling temperature used at every denoising step.
        commit_failure_budget: Maximum cumulative failure risk for normal commits.
        fused_entropy_weight: Multiplicative entropy penalty coefficient.
        jump_failure_budget: Maximum cumulative failure risk for forced jumps.
        vocab_chunk_size: Vocabulary projection planning size recorded with the
            model. Inference uses standard PyTorch tensor operations.
        latent_dim: Width of the recurrent latent state.
        latent_memory_slots: Number of persistent latent memory slots.
        latent_num_layers: Number of latent Transformer blocks.
        latent_num_heads: Number of latent attention heads.
        latent_local_attention_window: Local token-attention radius.
        latent_dropout: Latent Transformer dropout probability.
        jump_on_no_progress_after: Stagnation steps before a forced jump.
        max_ponder_steps: Maximum denoising iterations per requested token.
        min_trajectory_progress: Minimum fused-risk improvement counted as progress.
        turn_end_token_id: Native Gemma turn terminator.
        kwargs: Standard DiffusionGemma configuration values.
    """

    model_type = "modilify_mk1"
    sub_configs = {
        "text_config": ModilifyMk1TextConfig,
        **{
            key: value
            for key, value in DiffusionGemmaConfig.sub_configs.items()
            if key != "text_config"
        },
    }

    def __init__(
        self,
        text_config: (
            ModilifyMk1TextConfig
            | DiffusionGemmaTextConfig
            | dict[str, Any]
            | None
        ) = None,
        vision_config: Any | dict[str, Any] | None = None,
        *,
        denoise_temperature: float = 0.8,
        commit_failure_budget: float = 0.2,
        fused_entropy_weight: float = 0.5,
        jump_failure_budget: float = 2.0,
        vocab_chunk_size: int = 65_536,
        latent_dim: int = 1536,
        latent_memory_slots: int = 64,
        latent_num_layers: int = 4,
        latent_num_heads: int = 16,
        latent_local_attention_window: int = 128,
        latent_dropout: float = 0.0,
        jump_on_no_progress_after: int = 12,
        max_ponder_steps: int = 64,
        min_trajectory_progress: float = 0.005,
        turn_end_token_id: int = 106,
        **kwargs: Any,
    ) -> None:
        kwargs.pop("model_type", None)
        if isinstance(text_config, DiffusionGemmaTextConfig):
            text_payload = text_config.to_dict()
            text_payload.pop("model_type", None)
            text_config = ModilifyMk1TextConfig(**text_payload)
        elif isinstance(text_config, dict):
            text_payload = dict(text_config)
            text_payload.pop("model_type", None)
            text_config = ModilifyMk1TextConfig(**text_payload)
        elif text_config is None:
            text_config = ModilifyMk1TextConfig()

        self.denoise_temperature = float(denoise_temperature)
        self.commit_failure_budget = float(commit_failure_budget)
        self.fused_entropy_weight = float(fused_entropy_weight)
        self.jump_failure_budget = float(jump_failure_budget)
        self.vocab_chunk_size = int(vocab_chunk_size)
        self.latent_dim = int(latent_dim)
        self.latent_memory_slots = int(latent_memory_slots)
        self.latent_num_layers = int(latent_num_layers)
        self.latent_num_heads = int(latent_num_heads)
        self.latent_local_attention_window = int(latent_local_attention_window)
        self.latent_dropout = float(latent_dropout)
        self.jump_on_no_progress_after = int(jump_on_no_progress_after)
        self.max_ponder_steps = int(max_ponder_steps)
        self.min_trajectory_progress = float(min_trajectory_progress)
        self.turn_end_token_id = int(turn_end_token_id)
        super().__init__(
            text_config=text_config,
            vision_config=vision_config,
            **kwargs,
        )
        self.model_type = type(self).model_type
        if not hasattr(self, "eos_token_id"):
            self.eos_token_id = self.text_config.eos_token_id
        if not hasattr(self, "pad_token_id"):
            self.pad_token_id = self.text_config.pad_token_id
        if not hasattr(self, "bos_token_id"):
            self.bos_token_id = self.text_config.bos_token_id
        self._validate_modilify()

    def _validate_modilify(self) -> None:
        """Validate inference-only extension values."""

        policy_values = (
            self.denoise_temperature,
            self.commit_failure_budget,
            self.fused_entropy_weight,
            self.jump_failure_budget,
            self.min_trajectory_progress,
        )
        if any(not math.isfinite(value) for value in policy_values):
            raise ValueError("Modilify Mk1 policy values must be finite.")
        positive = (
            self.denoise_temperature,
            self.commit_failure_budget,
            self.jump_failure_budget,
            self.vocab_chunk_size,
            self.latent_dim,
            self.latent_memory_slots,
            self.latent_num_layers,
            self.latent_num_heads,
            self.latent_local_attention_window,
            self.jump_on_no_progress_after,
            self.max_ponder_steps,
        )
        if any(value <= 0 for value in positive):
            raise ValueError(
                "Modilify Mk1 dimensions, budgets, and intervals must be positive."
            )
        if self.fused_entropy_weight < 0:
            raise ValueError("`fused_entropy_weight` must be non-negative.")
        if self.latent_dim % self.latent_num_heads:
            raise ValueError("`latent_dim` must be divisible by `latent_num_heads`.")
        if not 0.0 <= self.latent_dropout < 1.0:
            raise ValueError("`latent_dropout` must be in [0, 1).")
        if self.min_trajectory_progress < 0:
            raise ValueError("`min_trajectory_progress` must be non-negative.")


__all__ = ["ModilifyMk1Config", "ModilifyMk1TextConfig"]