from __future__ import annotations

from typing import Any

from transformers import PretrainedConfig, Qwen3VLConfig, Qwen3VLTextConfig


class Qwen3VLStitchedConfig(PretrainedConfig):
    """Configuration for the experimental L16-to-L29 H3 conditioning frontend."""

    model_type = "qwen3_vl_stitched_h3"
    # A useful instance requires two complete heterogeneous backbone configs.
    # This also prevents PretrainedConfig.to_diff_dict() from constructing an
    # invalid no-argument comparison instance.
    has_no_defaults_at_init = True

    def __init__(
        self,
        *,
        source_config: dict[str, Any] | None = None,
        target_text_config: dict[str, Any] | None = None,
        bridge_config: dict[str, Any] | None = None,
        source_layer: int = 16,
        target_layer: int = 29,
        output_layer: int = 50,
        source_model_name: str = "Qwen/Qwen3-VL-8B-Instruct",
        target_model_name: str = "Qwen/Qwen3-VL-32B-Instruct",
        **kwargs: Any,
    ) -> None:
        kwargs.setdefault(
            "auto_map",
            {
                "AutoConfig": "configuration_qwen3_vl_stitched.Qwen3VLStitchedConfig",
                "AutoModel": (
                    "modeling_qwen3_vl_stitched.Qwen3VLStitchedForH3Conditioning"
                ),
            },
        )
        super().__init__(**kwargs)
        source = Qwen3VLConfig(**(source_config or {}))
        target = Qwen3VLTextConfig(**(target_text_config or {}))
        bridge = bridge_config or {
            "input_features": 4096,
            "output_features": 5120,
            "width": 512,
            "num_heads": 6,
            "head_dim": 64,
            "expand_v": 2.0,
            "mode": "chunk",
            "use_short_conv": True,
            "max_log_gain": 2.0,
            "num_lanes": 4,
            "output_head_reference_width": 512,
        }
        if not 0 < source_layer:
            raise ValueError("source_layer must be positive")
        if not 0 < target_layer < output_layer:
            raise ValueError("target_layer must be positive and precede output_layer")
        if source.text_config.hidden_size != int(bridge["input_features"]):
            raise ValueError("Source hidden size does not match the bridge input")
        if target.hidden_size != int(bridge["output_features"]):
            raise ValueError("Target hidden size does not match the bridge output")

        self.source_config = source.to_dict()
        self.target_text_config = target.to_dict()
        self.bridge_config = dict(bridge)
        self.source_layer = source_layer
        self.target_layer = target_layer
        self.output_layer = output_layer
        self.source_model_name = source_model_name
        self.target_model_name = target_model_name
        self.output_hidden_size = target.hidden_size
        self.is_encoder_decoder = False
        self.use_cache = False


__all__ = ["Qwen3VLStitchedConfig"]