"""WhisperMLA — Whisper with MLA (multi-head latent attention) converted decoder attention.

Standalone trust_remote_code modeling: builds the MLA architecture structurally at init (the
converted modules' weights then come from the checkpoint), so
``AutoModelForSpeechSeq2Seq.from_pretrained(repo, trust_remote_code=True)`` just works.
Pinned to transformers==4.46.x (the sibling vendored modeling targets it).

Adapted from Whisper-MLA / MHA2MLA (Apache-2.0); see the repo NOTICE.
"""
import types
from types import SimpleNamespace

from .configuration_whisper_mla import WhisperMLAConfig
from .modeling_whisper import WhisperAttention, WhisperForConditionalGeneration

# NOTE: patch_func is imported directly (not only transitively) so the HF dynamic-module loader
# copies it into the modules cache — it scans the ENTRY file's `from .x import y` lines.
from .patch_func import svd_low_rank_approx as _ensure_patch_func_is_bundled  # noqa: F401
from .patching_model_load import patch_model
from .patching_whisper import (
    custom_WhisperAttention_mla_forward,
    custom_WhisperAttention_nosplit_forward,
)

_SCOPE_TARGETS = {
    "decoder_self": frozenset({"decoder_self"}),
    "plus_cross": frozenset({"decoder_self", "decoder_cross"}),
    "all": frozenset({"decoder_self", "decoder_cross", "encoder_self"}),
}


def _structural_convert(model, config: WhisperMLAConfig) -> None:
    """Rebuild the attention modules as MLA per the config. Runs on the freshly-initialized model
    inside __init__ — the SVD of the (random) init weights only shapes the modules; real weights
    arrive from the checkpoint via from_pretrained's state-dict load."""
    if getattr(config, "num_key_value_heads", None) is None:
        config.num_key_value_heads = config.encoder_attention_heads
    preserve = config.mla_scheme == "partial_preservation"
    args = SimpleNamespace(
        rope_dim_for_mla=config.mla_rope_dim,
        partial_rope_version={"uniform": "uniform", "2_norm": "2-norm"}[config.mla_selection],
        low_rank=config.mla_low_rank,
        svd_init_method="joint",
        is_gqa2mha2mla=False,
        uniform_start_point=0,
        is_mla_from_scratch=False,
        qk_tensor_path=None,
        preserve_rope=preserve,
        targets=set(_SCOPE_TARGETS[config.mla_scope]),
    )
    patch_model(model, config, args)
    forward = custom_WhisperAttention_mla_forward if preserve else custom_WhisperAttention_nosplit_forward
    for module in model.modules():
        if isinstance(module, WhisperAttention) and hasattr(module, "kv_proj"):
            module.forward = types.MethodType(forward, module)


class WhisperMLAForConditionalGeneration(WhisperForConditionalGeneration):
    config_class = WhisperMLAConfig

    def __init__(self, config: WhisperMLAConfig):
        super().__init__(config)
        if config.mla_scheme != "none":
            _structural_convert(self, config)