"""Flint for Hugging Face transformers (remote code), ported from RecursiveSI's SGLang plugin.

Flint is Qwen3-MoE plus: learned attention sinks (``softmax_offset``), per-layer sliding or full
attention, no RoPE on ``no_rope_layer_indices``, sigmoid top-k routing with renormalized weights,
and an ungated shared expert in every MoE block. Experts are stored fused (transformers' Qwen3-MoE
layout); per-expert checkpoint tensors are merged on load and split again on save.

Portions adapted from transformers' Qwen3-MoE and GPT-OSS models (Apache License 2.0).
"""

from torch import nn
import torch
import torch.nn.functional as F
from transformers.cache_utils import Cache, DynamicCache
from transformers.conversion_mapping import (
    get_checkpoint_conversion_mapping,
    register_checkpoint_conversion_mapping,
)
from transformers.generation import GenerationMixin
from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
from transformers.modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
from transformers.modeling_utils import PreTrainedModel
from transformers.models.qwen3_moe.modeling_qwen3_moe import (
    Qwen3MoeExperts,
    Qwen3MoeMLP,
    Qwen3MoeRMSNorm,
    Qwen3MoeRotaryEmbedding,
    apply_rotary_pos_emb,
    repeat_kv,
)

from .configuration_flint import SLIDING_ATTENTION, FlintConfig

# Per-expert checkpoint tensors <-> fused experts, the same mapping transformers uses for Qwen2/3-MoE.
register_checkpoint_conversion_mapping("flint", get_checkpoint_conversion_mapping("qwen2_moe"), overwrite=True)


def eager_attention_with_sinks(module, query, key, value, attention_mask, scaling, dropout=0.0, **kwargs):
    """Softmax over [scores, sink]; the sink adds to the denominator and contributes no value."""
    key_states = repeat_kv(key, module.num_key_value_groups)
    value_states = repeat_kv(value, module.num_key_value_groups)
    scores = torch.matmul(query, key_states.transpose(2, 3)).float() * scaling
    if attention_mask is not None:
        scores = scores + attention_mask.float()
    sinks = module.softmax_offset.float().reshape(1, -1, 1, 1).expand(query.shape[0], -1, query.shape[-2], 1)
    logits = torch.cat([scores, sinks], dim=-1)
    probs = F.softmax(logits - logits.amax(dim=-1, keepdim=True), dim=-1)[..., :-1]
    attn_output = torch.matmul(probs.to(value_states.dtype), value_states)
    return attn_output.transpose(1, 2).contiguous(), probs


class FlintAttention(nn.Module):
    """Grouped-query attention with QK RMSNorm, learned sinks, optional sliding window and RoPE."""

    def __init__(self, config: FlintConfig, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.head_dim = config.head_dim
        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
        self.scaling = self.head_dim**-0.5
        self.is_causal = True
        self.skip_rope = layer_idx in set(config.no_rope_layer_indices)
        bias = config.attention_bias
        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=bias)
        self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=bias)
        self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=bias)
        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=bias)
        self.q_norm = Qwen3MoeRMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.k_norm = Qwen3MoeRMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.softmax_offset = nn.Parameter(torch.zeros(config.num_attention_heads))

    def forward(self, hidden_states, position_embeddings, attention_mask, past_key_values=None, **kwargs):
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)
        query = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
        key = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
        value = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        if not self.skip_rope:
            cos, sin = position_embeddings
            query, key = apply_rotary_pos_emb(query, key, cos, sin)
        if past_key_values is not None:
            key, value = past_key_values.update(key, value, self.layer_idx)
        attn_output, attn_weights = eager_attention_with_sinks(self, query, key, value, attention_mask, self.scaling)
        return self.o_proj(attn_output.reshape(*input_shape, -1)), attn_weights


class FlintRouter(nn.Module):
    """Sigmoid scores, top-k, weights renormalized to sum to one."""

    # shortcut: no expert_bias (Flint's checkpoint has none; SGLang leaves it at zero). Add it if a checkpoint ships one.

    def __init__(self, config: FlintConfig):
        super().__init__()
        self.top_k = config.num_experts_per_tok
        self.norm_topk_prob = config.norm_topk_prob
        self.weight = nn.Parameter(torch.zeros(config.num_experts, config.hidden_size))

    def forward(self, hidden_states):
        router_logits = F.linear(hidden_states, self.weight)
        scores = torch.sigmoid(router_logits.float())
        top_weights, top_ids = torch.topk(scores, self.top_k, dim=-1)
        if self.norm_topk_prob:
            top_weights = top_weights / (top_weights.sum(dim=-1, keepdim=True) + 1e-20)
        return router_logits, top_weights.to(hidden_states.dtype), top_ids


class FlintSparseMoeBlock(nn.Module):
    def __init__(self, config: FlintConfig):
        super().__init__()
        self.experts = Qwen3MoeExperts(config)
        self.gate = FlintRouter(config)
        self.shared_expert = Qwen3MoeMLP(config, intermediate_size=config.shared_expert_intermediate_size)

    def forward(self, hidden_states):
        batch, seq, dim = hidden_states.shape
        flat = hidden_states.view(-1, dim)
        _, weights, ids = self.gate(flat)
        out = self.experts(flat, ids, weights) + self.shared_expert(flat)
        return out.view(batch, seq, dim)


class FlintDecoderLayer(nn.Module):
    def __init__(self, config: FlintConfig, layer_idx: int):
        super().__init__()
        self.self_attn = FlintAttention(config, layer_idx)
        if layer_idx in config.mlp_only_layers:
            self.mlp = Qwen3MoeMLP(config, intermediate_size=config.intermediate_size)
        else:
            self.mlp = FlintSparseMoeBlock(config)
        self.input_layernorm = Qwen3MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = Qwen3MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(self, hidden_states, attention_mask=None, position_embeddings=None, past_key_values=None, **kwargs):
        residual = hidden_states
        hidden_states, _ = self.self_attn(
            self.input_layernorm(hidden_states), position_embeddings, attention_mask, past_key_values
        )
        hidden_states = residual + hidden_states
        return hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))


class FlintPreTrainedModel(PreTrainedModel):
    config_class = FlintConfig
    config: FlintConfig
    base_model_prefix = "model"
    _no_split_modules = ["FlintDecoderLayer"]
    _skip_keys_device_placement = ["past_key_values"]
    _supports_flash_attn = False
    _supports_sdpa = False
    _supports_flex_attn = False
    # Experts are the decorated Qwen3MoeExperts, but transformers looks for the decorator in this file's source.
    _can_set_experts_implementation_cached_value = True

    @torch.no_grad()
    def _init_weights(self, module):
        super()._init_weights(module)
        std = self.config.initializer_range
        if isinstance(module, Qwen3MoeExperts):
            nn.init.normal_(module.gate_up_proj, mean=0.0, std=std)
            nn.init.normal_(module.down_proj, mean=0.0, std=std)
        elif isinstance(module, FlintRouter):
            nn.init.normal_(module.weight, mean=0.0, std=std)
        elif isinstance(module, FlintAttention):
            nn.init.zeros_(module.softmax_offset)


class FlintModel(FlintPreTrainedModel):
    def __init__(self, config: FlintConfig):
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
        self.layers = nn.ModuleList([FlintDecoderLayer(config, i) for i in range(config.num_hidden_layers)])
        self.norm = Qwen3MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.rotary_emb = Qwen3MoeRotaryEmbedding(config=config)
        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, value):
        self.embed_tokens = value

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        position_ids=None,
        past_key_values: Cache | None = None,
        inputs_embeds=None,
        use_cache=None,
        output_hidden_states=None,
        **kwargs,
    ) -> MoeModelOutputWithPast:
        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("Pass exactly one of input_ids or inputs_embeds")
        use_cache = self.config.use_cache if use_cache is None else use_cache
        output_hidden_states = (
            self.config.output_hidden_states if output_hidden_states is None else output_hidden_states
        )
        if use_cache and past_key_values is None:
            past_key_values = DynamicCache(config=self.config)
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)
        if position_ids is None:
            past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
            position_ids = (torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen).unsqueeze(0)

        if not isinstance(mask_mapping := attention_mask, dict):
            mask_kwargs = {
                "config": self.config,
                "inputs_embeds": inputs_embeds,
                "attention_mask": attention_mask,
                "past_key_values": past_key_values,
                "position_ids": position_ids,
            }
            mask_mapping = {
                "full_attention": create_causal_mask(**mask_kwargs),
                SLIDING_ATTENTION: create_sliding_window_causal_mask(**mask_kwargs),
            }

        hidden_states = inputs_embeds
        position_embeddings = self.rotary_emb(hidden_states, position_ids)
        all_hidden = (hidden_states,) if output_hidden_states else None
        for i, layer in enumerate(self.layers):
            hidden_states = layer(
                hidden_states,
                attention_mask=mask_mapping[self.config.layer_types[i]],
                position_embeddings=position_embeddings,
                past_key_values=past_key_values,
            )
            if output_hidden_states:
                all_hidden += (hidden_states,)
        hidden_states = self.norm(hidden_states)
        if output_hidden_states:
            # Match transformers' convention: the last entry is the normed output.
            all_hidden = all_hidden[:-1] + (hidden_states,)
        return MoeModelOutputWithPast(
            last_hidden_state=hidden_states, past_key_values=past_key_values, hidden_states=all_hidden
        )


class FlintForCausalLM(FlintPreTrainedModel, GenerationMixin):
    _tied_weights_keys = {}

    def __init__(self, config: FlintConfig):
        super().__init__(config)
        self.model = FlintModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def get_output_embeddings(self):
        return self.lm_head

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        position_ids=None,
        past_key_values=None,
        inputs_embeds=None,
        labels=None,
        use_cache=None,
        output_hidden_states=None,
        logits_to_keep=0,
        **kwargs,
    ) -> MoeCausalLMOutputWithPast:
        outputs = 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,
            output_hidden_states=output_hidden_states,
        )
        hidden_states = outputs.last_hidden_state
        keep = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
        logits = self.lm_head(hidden_states[:, keep, :])
        loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) if labels is not None else None
        return MoeCausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
        )


__all__ = ["FlintForCausalLM", "FlintModel", "FlintPreTrainedModel"]