"""Configuration for the SPEV multimodal decision model."""

import math
from typing import ClassVar

from transformers import PretrainedConfig, Qwen3OmniMoeThinkerConfig


class SPEVConfig(PretrainedConfig):
    model_type = "spev"
    sub_configs: ClassVar[dict] = {"thinker_config": Qwen3OmniMoeThinkerConfig}

    def __init__(
        self,
        thinker_config=None,
        adapter_rank=32,
        adapter_alpha=64,
        head_width=256,
        head_heads=8,
        head_ff=1024,
        head_depth=2,
        head_architecture="isolated_evidence_v1",
        temperature=1.0,
        max_tokens=8192,
        **kwargs,
    ):
        if isinstance(thinker_config, dict):
            thinker_config = Qwen3OmniMoeThinkerConfig(**thinker_config)
        self.thinker_config = thinker_config or Qwen3OmniMoeThinkerConfig()
        for name, value in (
            ("adapter_rank", adapter_rank),
            ("head_width", head_width),
            ("head_heads", head_heads),
            ("head_ff", head_ff),
            ("head_depth", head_depth),
            ("max_tokens", max_tokens),
        ):
            if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
                raise ValueError(f"{name} must be a positive integer")
        if not math.isfinite(adapter_alpha) or adapter_alpha <= 0:
            raise ValueError("adapter_alpha must be positive and finite")
        if not math.isfinite(temperature) or temperature <= 0:
            raise ValueError("temperature must be positive and finite")
        self.adapter_rank = adapter_rank
        self.adapter_alpha = adapter_alpha
        self.head_width = head_width
        if head_architecture != "isolated_evidence_v1":
            raise ValueError("Unsupported SPEV decision head architecture")
        if head_width % head_heads:
            raise ValueError("head_width must be divisible by head_heads")
        self.head_heads = head_heads
        self.head_ff = head_ff
        self.head_depth = head_depth
        self.head_architecture = head_architecture
        self.temperature = temperature
        self.max_tokens = max_tokens
        kwargs.setdefault("tie_word_embeddings", False)
        super().__init__(**kwargs)

    def get_text_config(self, decoder=False, encoder=False):
        return self.thinker_config.text_config


SPEVConfig.register_for_auto_class()