"""MinistralDualRope — Ministral with a per-layer sliding/full attention pattern AND a second,
unscaled rotary for the sliding layers. Loadable via ``trust_remote_code``.

This is the MINIMAL delta over the *stock* transformers Ministral model:

  * Ministral already provides the Llama-style math (RMSNorm without unit offset,
    SiLU-gated MLP, GQA, no biases, no QK-norm, no sandwich norms, no softcapping) AND
    the per-layer sliding/full attention pattern routed by ``config.layer_types`` — the
    attention/mask routing is bit-identical to how this model was trained.
  * The ONE thing Ministral lacks is a *second* RoPE. Our full-attention layers use the
    (llama3-scaled) global rotary; the sliding layers use an UNSCALED rotary at
    ``rope_local_base_freq`` (the Gemma-3 dual-RoPE design). We add ``rotary_emb_local``
    and route ``position_embeddings`` per ``layer_types``. Everything else is stock
    Ministral, inherited unchanged.

The building blocks are pulled from the *installed* ``modeling_ministral`` so this stays
aligned with whatever transformers version loads the checkpoint. Verified to match the
faithful FlexLlama reference to the analytic-vs-training-buffer RoPE floor (FP32 top-1
agreement ~100%, mean |Δ logit| ~2e-3).
"""

from __future__ import annotations

import copy
import inspect
from functools import partial
from typing import Optional

import torch

from transformers.models.ministral import modeling_ministral as _M

from .configuration_ministral_dual_rope import MinistralDualRopeConfig

MinistralModel = _M.MinistralModel
MinistralForCausalLM = _M.MinistralForCausalLM
MinistralRotaryEmbedding = _M.MinistralRotaryEmbedding
create_causal_mask = _M.create_causal_mask
create_sliding_window_causal_mask = _M.create_sliding_window_causal_mask
DynamicCache = _M.DynamicCache
BaseModelOutputWithPast = _M.BaseModelOutputWithPast
# `check_model_inputs` moved between transformers minor/major versions; fall back to a
# no-op so the forward still works if a given install lacks it. We deliberately do NOT
# apply transformers' `auto_docstring` to the forward: on tf>=5 it scans the signature at
# import and logs a scary `[ERROR] ... is part of ...'s signature, but not documented` for
# every undocumented kwarg (e.g. `cache_position`). That is pure lint noise — the forward
# is fully functional without it — but it panics downstream users, so we omit it.
check_model_inputs = getattr(_M, "check_model_inputs", lambda f: f)

# The mask-builder signature drifts across transformers versions (the embeds kwarg was
# renamed ``input_embeds`` -> ``inputs_embeds`` and ``cache_position`` was dropped between
# 4.x and 5.x). Introspect once and pass only what each installed builder accepts, so the
# same shipped file loads under both. Computed at import time (cheap, version-stable).
_CAUSAL_MASK_PARAMS = set(inspect.signature(create_causal_mask).parameters)
_SLIDING_MASK_PARAMS = set(inspect.signature(create_sliding_window_causal_mask).parameters)


class MinistralDualRopeModel(MinistralModel):
    config_class = MinistralDualRopeConfig

    def __init__(self, config: MinistralDualRopeConfig):
        super().__init__(config)
        # Second, UNSCALED rotary for the sliding (local) layers. The parent's
        # ``self.rotary_emb`` is the (possibly llama3-scaled) GLOBAL rotary used by the
        # full-attention layers; build a sibling at ``rope_local_base_freq`` with NO
        # scaling. When base == rope_theta and there is no scaling the two are identical.
        local_cfg = copy.deepcopy(config)
        base = float(getattr(config, "rope_local_base_freq", None) or config._global_rope_theta())
        # Configure an UNSCALED rotary. The two transformers major versions expose RoPE
        # differently and the branches MUST be mutually exclusive:
        #   * transformers >= 5 stores everything in the unified ``rope_parameters`` dict
        #     and exposes ``rope_scaling`` as an ALIAS *onto* it — so setting
        #     ``rope_scaling`` here would clobber ``rope_parameters`` back to ``None`` and
        #     blow up inside ``MinistralRotaryEmbedding``. Set ONLY ``rope_parameters``.
        #   * transformers 4.x has no ``rope_parameters``; the rotary reads ``rope_scaling``
        #     (must be None to disable llama3 scaling) + a top-level ``rope_theta``.
        if hasattr(local_cfg, "rope_parameters"):        # transformers >= 5
            local_cfg.rope_parameters = {"rope_type": "default", "rope_theta": base}
        else:                                            # transformers 4.x
            local_cfg.rope_scaling = None
            local_cfg.rope_theta = base
        self.rotary_emb_local = MinistralRotaryEmbedding(config=local_cfg)
        self.post_init()

    @check_model_inputs
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values=None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        **kwargs,
    ) -> BaseModelOutputWithPast:
        # Verbatim from MinistralModel.forward except the two dual-RoPE lines marked (*).
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        if use_cache and past_key_values is None:
            past_key_values = DynamicCache(config=self.config)

        if cache_position is None:
            past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
            cache_position = torch.arange(
                past_seen, past_seen + inputs_embeds.shape[1], device=inputs_embeds.device
            )

        if position_ids is None:
            position_ids = cache_position.unsqueeze(0)

        # `generate` may already pass a prepared {type: mask} dict; else build both masks.
        if not isinstance(causal_mask_mapping := attention_mask, dict):
            # Superset of every kwarg any transformers version's mask builder wants; the
            # two embeds spellings and cache_position are filtered per the installed
            # signature (see _CAUSAL_MASK_PARAMS above) so this works on 4.x and >=5.
            _mask_src = {
                "config": self.config,
                "input_embeds": inputs_embeds,   # transformers 4.x
                "inputs_embeds": inputs_embeds,  # transformers >= 5
                "attention_mask": attention_mask,
                "cache_position": cache_position,  # transformers 4.x only
                "past_key_values": past_key_values,
                "position_ids": position_ids,
            }
            causal_mask_mapping = {
                "full_attention": create_causal_mask(
                    **{k: v for k, v in _mask_src.items() if k in _CAUSAL_MASK_PARAMS}),
                "sliding_attention": create_sliding_window_causal_mask(
                    **{k: v for k, v in _mask_src.items() if k in _SLIDING_MASK_PARAMS}),
            }

        hidden_states = inputs_embeds
        # (*) Dual RoPE: full-attention layers get the (scaled) global rotary, sliding
        #     layers get the unscaled local rotary — routed by the same layer type as the
        #     mask above.
        position_embeddings = {
            "full_attention": self.rotary_emb(hidden_states, position_ids),
            "sliding_attention": self.rotary_emb_local(hidden_states, position_ids),
        }

        for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
            # (*) Route the mask AND the rope by the layer's attention type. Read it from
            #     ``config.layer_types`` (stable across versions) rather than the layer
            #     attribute, which transformers renamed ``attention_type`` -> ``layer_type``
            #     between 4.x and 5.x.
            layer_type = self.config.layer_types[i]
            layer_kwargs = dict(
                attention_mask=causal_mask_mapping[layer_type],
                position_ids=position_ids,
                past_key_values=past_key_values,
                use_cache=use_cache,
                cache_position=cache_position,
                position_embeddings=position_embeddings[layer_type],
                **kwargs,
            )
            if self.gradient_checkpointing and self.training:
                # Honor activation checkpointing. The stock MinistralModel.forward does
                # this; our override MUST replicate it or long-context SFT/RL (this is a
                # 32K model) OOMs — even under LoRA, since activations, not optimizer
                # state, dominate at long sequence length. Bake the per-layer kwargs into
                # a partial so we do NOT depend on the decoder layer's positional arg
                # order, which drifts across transformers versions.
                hidden_states = self._gradient_checkpointing_func(
                    partial(decoder_layer.__call__, **layer_kwargs), hidden_states
                )
            else:
                hidden_states = decoder_layer(hidden_states, **layer_kwargs)

        hidden_states = self.norm(hidden_states)
        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=past_key_values if use_cache else None,
        )


class MinistralDualRopeForCausalLM(MinistralForCausalLM):
    config_class = MinistralDualRopeConfig

    def __init__(self, config: MinistralDualRopeConfig):
        super().__init__(config)
        self.model = MinistralDualRopeModel(config)
        self.post_init()


__all__ = ["MinistralDualRopeForCausalLM", "MinistralDualRopeModel"]