# Copyright 2026 The OpenBMB Team and the HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import math
from collections.abc import Callable

import torch
import torch.nn.functional as F
from huggingface_hub.dataclasses import strict

from ... import initialization as init
from ...cache_utils import Cache
from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
from ...processing_utils import Unpack
from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
from ...utils.generic import is_flash_attention_requested
from ..deepseek_v2.modeling_deepseek_v2 import DeepseekV2Attention
from ..gemma3.modeling_gemma3 import Gemma3TextScaledWordEmbedding
from ..llama.configuration_llama import LlamaConfig
from ..llama.modeling_llama import (
    LlamaDecoderLayer,
    LlamaForCausalLM,
    LlamaForSequenceClassification,
    LlamaMLP,
    LlamaModel,
    LlamaPreTrainedModel,
    LlamaRMSNorm,
    LlamaRotaryEmbedding,
    apply_rotary_pos_emb,
    eager_attention_forward,
)


logger = logging.get_logger(__name__)


@auto_docstring(checkpoint="openbmb/MiniCPM3-4B")
@strict
class MiniCPM3Config(LlamaConfig):
    r"""
    kv_lora_rank (`int`, *optional*, defaults to 256):
        Rank of the low-rank KV projection in multi-head latent attention.
    q_lora_rank (`int`, *optional*, defaults to 768):
        Rank of the low-rank query projection in multi-head latent attention. If `None`, the query projection
        is a single dense projection rather than a low-rank one.
    qk_nope_head_dim (`int`, *optional*, defaults to 64):
        Dimension of the non-RoPE part of each query/key head.
    qk_rope_head_dim (`int`, *optional*, defaults to 32):
        Dimension of the RoPE part of each query/key head.
    v_head_dim (`int`, *optional*):
        Dimension of each value head. If `None`, defaults to `hidden_size // num_attention_heads`.
    scale_emb (`int` or `float`, *optional*, defaults to 12):
        Multiplier applied to input embeddings.
    scale_depth (`int` or `float`, *optional*, defaults to 1.4):
        Multiplier for residual connections; the effective scaling is `scale_depth / sqrt(num_hidden_layers)`.
        If `None`, defaults to `sqrt(num_hidden_layers)` (no-op scaling).
    dim_model_base (`int`, *optional*, defaults to 256):
        Base model dimension used to scale logits before the language model head. If `None`,
        defaults to `hidden_size` (no-op scaling).

    Example:

    ```python
    >>> from transformers import MiniCPM3Model, MiniCPM3Config
    >>> # Initializing a MiniCPM3 style configuration
    >>> configuration = MiniCPM3Config()
    >>> # Initializing a model from the configuration
    >>> model = MiniCPM3Model(configuration)
    >>> # Accessing the model configuration
    >>> configuration = model.config
    ```
    """

    base_model_tp_plan = {
        "layers.*.self_attn.q_proj": "colwise",
        "layers.*.self_attn.q_b_proj": "colwise",
        "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
        "layers.*.self_attn.kv_b_proj": "colwise",
        "layers.*.self_attn.o_proj": "rowwise",
        "layers.*.mlp.gate_proj": "colwise",
        "layers.*.mlp.up_proj": "colwise",
        "layers.*.mlp.down_proj": "rowwise",
    }

    model_type = "minicpm3"

    # Only fields whose defaults differ from `LlamaConfig` are redeclared here; the rest are inherited.
    # Defaults match the `openbmb/MiniCPM3-4B` checkpoint.
    vocab_size: int = 73448
    hidden_size: int = 2560
    intermediate_size: int = 6400
    num_hidden_layers: int = 62
    num_attention_heads: int = 40
    num_key_value_heads: int | None = 40
    max_position_embeddings: int = 32768
    initializer_range: float = 0.1
    rms_norm_eps: float = 1e-5
    tie_word_embeddings: bool = True
    kv_lora_rank: int = 256
    q_lora_rank: int | None = 768
    qk_nope_head_dim: int = 64
    qk_rope_head_dim: int = 32
    v_head_dim: int | None = None
    scale_emb: int | float = 12
    scale_depth: int | float | None = 1.4
    dim_model_base: int | None = 256

    def __post_init__(self, **kwargs):
        # In MLA the per-head dim used by RoPE is the rotary part, not `hidden_size / num_attention_heads`.
        self.head_dim = self.qk_rope_head_dim
        # When explicitly set to `None`, these collapse to no-op scalings (e.g. for a randomly
        # initialised tiny model); the class defaults otherwise match `openbmb/MiniCPM3-4B`.
        if self.v_head_dim is None:
            self.v_head_dim = self.hidden_size // self.num_attention_heads
        if self.scale_depth is None:
            self.scale_depth = math.sqrt(self.num_hidden_layers)
        if self.dim_model_base is None:
            self.dim_model_base = self.hidden_size
        super().__post_init__(**kwargs)

    @property
    def logits_scaling(self) -> float:
        # Hidden states are divided by this factor before the LM head (`1` when `dim_model_base == hidden_size`).
        return self.hidden_size / self.dim_model_base


class MiniCPM3ScaledWordEmbedding(Gemma3TextScaledWordEmbedding):
    pass


class MiniCPM3RMSNorm(LlamaRMSNorm):
    pass


class MiniCPM3RotaryEmbedding(LlamaRotaryEmbedding):
    pass


class MiniCPM3MLP(LlamaMLP):
    pass


class MiniCPM3Attention(DeepseekV2Attention):
    """
    Multi-head Latent Attention (MLA), structurally identical to `DeepseekV2Attention`.
    The only difference is the rotary convention: MiniCPM3 keeps the original cos/sin RoPE
    (`apply_rotary_pos_emb`) instead of DeepSeek-V2's complex rotary, so we inherit the
    module construction and override only `forward`.
    """

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: torch.Tensor | None = None,
        past_key_values: Cache | None = None,
        **kwargs,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        batch_size, seq_length = hidden_states.shape[:-1]
        query_shape = (batch_size, seq_length, -1, self.qk_head_dim)
        key_shape = (batch_size, seq_length, -1, self.qk_nope_head_dim + self.v_head_dim)

        if self.q_lora_rank is None:
            q_states = self.q_proj(hidden_states)
        else:
            q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))
        q_states = q_states.view(query_shape).transpose(1, 2)
        q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

        compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
        k_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)

        k_pass = self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(key_shape).transpose(1, 2)
        k_pass, value_states = torch.split(k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)

        k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim)

        cos, sin = position_embeddings
        # Same MLA forward as DeepSeek-V2/V3, except MiniCPM3 keeps the standard (non-interleaved)
        # cos/sin rotary (`apply_rotary_pos_emb`) instead of DeepSeek's complex/interleaved variant.
        q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin)
        k_rot = k_rot.expand(*k_pass.shape[:-1], -1)

        query_states = torch.cat((q_pass, q_rot), dim=-1)
        key_states = torch.cat((k_pass, k_rot), dim=-1)

        if past_key_values is not None:
            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)

        if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim:
            value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim])

        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
            self.config._attn_implementation, eager_attention_forward
        )

        attn_output, attn_weights = attention_interface(
            self,
            query_states,
            key_states,
            value_states,
            attention_mask,
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
            **kwargs,
        )

        if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim:
            attn_output = attn_output[:, :, :, : self.v_head_dim]

        attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous()
        attn_output = self.o_proj(attn_output)
        return attn_output, attn_weights


class MiniCPM3DecoderLayer(LlamaDecoderLayer):
    def __init__(self, config: MiniCPM3Config, layer_idx: int):
        super().__init__(config, layer_idx)
        # MiniCPM3 multiplies each residual branch by `scale_depth / sqrt(num_hidden_layers)`
        # (Llama adds the branch directly). Precompute the constant once instead of per forward.
        self.residual_scale = config.scale_depth / math.sqrt(config.num_hidden_layers)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        use_cache: bool | None = False,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> torch.Tensor:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states, _ = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            position_embeddings=position_embeddings,
            **kwargs,
        )
        hidden_states = residual + hidden_states * self.residual_scale

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states * self.residual_scale
        return hidden_states


class MiniCPM3PreTrainedModel(LlamaPreTrainedModel):
    @torch.no_grad()
    def _init_weights(self, module):
        PreTrainedModel._init_weights(self, module)
        if isinstance(module, MiniCPM3ScaledWordEmbedding):
            init.constant_(module.embed_scale, module.scalar_embed_scale)


@auto_docstring
class MiniCPM3Model(LlamaModel):
    def __init__(self, config: MiniCPM3Config):
        super().__init__(config)
        # MiniCPM3 scales the input embeddings by `scale_emb` (not present in Llama).
        self.embed_tokens = MiniCPM3ScaledWordEmbedding(
            config.vocab_size, config.hidden_size, self.padding_idx, embed_scale=config.scale_emb
        )


@auto_docstring
class MiniCPM3ForCausalLM(LlamaForCausalLM):
    @can_return_tuple
    @auto_docstring
    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        labels: torch.LongTensor | None = None,
        use_cache: bool | None = None,
        logits_to_keep: int | torch.Tensor = 0,
        **kwargs: Unpack[TransformersKwargs],
    ) -> CausalLMOutputWithPast:
        r"""
        Example:

        ```python
        >>> from transformers import AutoTokenizer, MiniCPM3ForCausalLM

        >>> model = MiniCPM3ForCausalLM.from_pretrained("openbmb/MiniCPM3-4B")
        >>> tokenizer = AutoTokenizer.from_pretrained("openbmb/MiniCPM3-4B")

        >>> prompt = "Hey, are you conscious? Can you talk to me?"
        >>> inputs = tokenizer(prompt, return_tensors="pt")

        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
        ```"""
        outputs: BaseModelOutputWithPast = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            **kwargs,
        )

        hidden_states = outputs.last_hidden_state
        # MiniCPM3 scales hidden states down before the LM head (not present in Llama).
        hidden_states = hidden_states / self.config.logits_scaling
        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
        logits = self.lm_head(hidden_states[:, slice_indices, :])

        loss = None
        if labels is not None:
            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


class MiniCPM3ForSequenceClassification(LlamaForSequenceClassification):
    pass


__all__ = [
    "MiniCPM3Config",
    "MiniCPM3PreTrainedModel",
    "MiniCPM3Model",
    "MiniCPM3ForCausalLM",
    "MiniCPM3ForSequenceClassification",
]