import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers.models.lfm2_moe.modeling_lfm2_moe import (
    Lfm2MoeForCausalLM,
    Lfm2MoeSparseMoeBlock,
    Lfm2MoeTopKRouter,
    Lfm2MoeExperts,
    Lfm2MoeModel,
)
from transformers.cache_utils import DynamicCache
from transformers.masking_utils import create_causal_mask, create_recurrent_attention_mask
from transformers.modeling_outputs import MoeModelOutputWithPast
from .configuration_lfm2_moe_custom import Lfm2MoeCustomConfig

# ===================================================================
# HOTFIX: Installed transformers is missing "linear_attention" key
# in causal_mask_mapping. Replace the base forward at runtime.
# ===================================================================
_original_lfm_forward = Lfm2MoeModel.forward

def _lfm_forward_fixed(
    self,
    input_ids: torch.LongTensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.LongTensor | None = None,
    past_key_values=None,
    inputs_embeds: torch.FloatTensor | None = None,
    use_cache: bool | None = None,
    **kwargs,
):
    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 position_ids is None:
        past_seen_tokens = 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_tokens
        position_ids = position_ids.unsqueeze(0)

    if not isinstance(causal_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,
        }
        causal_mask_mapping = {
            "full_attention": create_causal_mask(**mask_kwargs),
            "conv": create_recurrent_attention_mask(**mask_kwargs),
            "linear_attention": None,
        }

    hidden_states = inputs_embeds
    position_embeddings = self.pos_emb(hidden_states, position_ids=position_ids)

    for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
        hidden_states = decoder_layer(
            hidden_states,
            attention_mask=causal_mask_mapping[self.config.layer_types[i]],
            position_ids=position_ids,
            past_key_values=past_key_values,
            position_embeddings=position_embeddings,
            **kwargs,
        )

    hidden_states = self.embedding_norm(hidden_states)

    return MoeModelOutputWithPast(
        last_hidden_state=hidden_states,
        past_key_values=past_key_values,
    )

Lfm2MoeModel.forward = _lfm_forward_fixed
# ===================================================================

class Lfm2MoeDynamicTopKRouter(Lfm2MoeTopKRouter):
    def __init__(self, config: Lfm2MoeCustomConfig, layer_idx: int = -1):
        nn.Module.__init__(self)
        self.num_experts = config.num_experts
        self.norm_topk_prob = config.norm_topk_prob
        self.hidden_dim = config.hidden_size
        self.weight = nn.Parameter(torch.zeros(self.num_experts, self.hidden_dim))
        self.routed_scaling_factor = config.routed_scaling_factor
        self.use_expert_bias = config.use_expert_bias

        self.budget = config.budget
        self.min_dynamic_k = 4
        self.max_dynamic_k = min(4, self.num_experts)
        self.layer_idx = layer_idx

        # ---- EXPERT ACTIVATION STATISTICS ----
        self.collect_stats = False
        self._stats_sum = 0.0      # sum of dynamic_k values
        self._stats_sum_sq = 0.0   # sum of squared dynamic_k values
        self._stats_count = 0      # total tokens seen
        # --------------------------------------

    @staticmethod
    def inverse_simpson_ratio(probs, dim=-1, eps=1e-12):
        """ISR = 1 / sum(p_i^2). Measures effective number of choices."""
        return 1.0 / (probs.pow(2).sum(dim=dim) + eps)

    def forward(self, hidden_states, expert_bias=None):
        router_logits = F.linear(hidden_states, self.weight)

        # ---- DYNAMIC TOP-K via ISR ----
        std = router_logits.std(dim=-1, keepdim=True)
        # norm_logits = router_logits / (std + 1e-12)
        norm_logits = router_logits
        probs = F.softmax(norm_logits, dim=-1)
        isr = self.inverse_simpson_ratio(probs, dim=-1)
        dynamic_k = torch.ceil(isr * self.budget).long().clamp(
            self.min_dynamic_k, self.max_dynamic_k
        )
        # -------------------------------

        # ---- RECORD STATS (zero overhead when disabled) ----
        if self.collect_stats:
            # dynamic_k shape: [batch, seq] or [batch*seq]
            k_vals = dynamic_k.detach().cpu().float()
            self._stats_sum += k_vals.sum().item()
            self._stats_sum_sq += (k_vals ** 2).sum().item()
            self._stats_count += k_vals.numel()
        # ----------------------------------------------------

        routing_weights = router_logits.sigmoid()

        batch_max_k = int(dynamic_k.max().item())
        batch_max_k = max(batch_max_k, self.min_dynamic_k)
        batch_max_k = min(batch_max_k, self.max_dynamic_k)

        if self.use_expert_bias:
            scores_for_routing = routing_weights + expert_bias
            _, selected_experts = torch.topk(scores_for_routing, k=batch_max_k, dim=-1)
            gathered_weights = torch.gather(
                routing_weights, dim=1, index=selected_experts
            ).type_as(router_logits)
        else:
            gathered_weights, selected_experts = torch.topk(
                routing_weights, k=batch_max_k, dim=-1
            )

        positions = torch.arange(batch_max_k, device=hidden_states.device).unsqueeze(0)
        valid_mask = positions < dynamic_k.unsqueeze(1)
        dummy_idx = self.num_experts

        selected_experts = torch.where(
            valid_mask, selected_experts, torch.full_like(selected_experts, dummy_idx)
        )
        gathered_weights = torch.where(
            valid_mask, gathered_weights, torch.zeros_like(gathered_weights)
        )

        if self.norm_topk_prob:
            sum_weights = gathered_weights.sum(dim=-1, keepdim=True)
            gathered_weights = gathered_weights / (sum_weights + 1e-6)

        routing_weights = gathered_weights * self.routed_scaling_factor
        return router_logits, routing_weights, selected_experts


class Lfm2MoeDynamicExperts(Lfm2MoeExperts):
    """
    Same as Lfm2MoeExperts but skips dummy expert indices (== num_experts)
    so that masked-out top-k slots do not trigger wasted MLP compute.
    """

    def __init__(self, config):
        super().__init__(config)

    def forward(self, hidden_states, top_k_index, top_k_weights):
        final_hidden_states = torch.zeros_like(hidden_states)
        dummy = self.num_experts

        with torch.no_grad():
            expert_mask = torch.nn.functional.one_hot(
                top_k_index, num_classes=self.num_experts + 1
            )
            expert_mask = expert_mask.permute(2, 1, 0)
            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()

        for expert_idx in expert_hit:
            expert_idx = expert_idx[0]
            if expert_idx == dummy:
                continue
            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
            current_state = hidden_states[token_idx]
            gate, up = nn.functional.linear(
                current_state, self.gate_up_proj[expert_idx]
            ).chunk(2, dim=-1)
            current_hidden_states = self.act_fn(gate) * up
            current_hidden_states = nn.functional.linear(
                current_hidden_states, self.down_proj[expert_idx]
            )
            current_hidden_states = current_hidden_states * top_k_weights[
                token_idx, top_k_pos, None
            ]
            final_hidden_states.index_add_(
                0, token_idx, current_hidden_states.to(final_hidden_states.dtype)
            )

        return final_hidden_states


class Lfm2MoeCustomSparseMoeBlock(Lfm2MoeSparseMoeBlock):
    def __init__(self, config: Lfm2MoeCustomConfig, layer_idx: int):
        nn.Module.__init__(self)
        self.experts = Lfm2MoeDynamicExperts(config)
        self.gate = Lfm2MoeDynamicTopKRouter(config, layer_idx=layer_idx)
        self.use_expert_bias = config.use_expert_bias
        if self.use_expert_bias:
            self.register_buffer(
                "expert_bias", torch.zeros(config.num_experts, dtype=torch.float32)
            )


class Lfm2MoeCustomForCausalLM(Lfm2MoeForCausalLM):
    config_class = Lfm2MoeCustomConfig

    def __init__(self, config: Lfm2MoeCustomConfig):
        super().__init__(config)
        for layer_idx, layer in enumerate(self.model.layers):
            if layer_idx >= config.num_dense_layers:
                layer.feed_forward = Lfm2MoeCustomSparseMoeBlock(config, layer_idx)

    # ------------------------------------------------------------------
    # EXPERT ACTIVATION STATISTICS API
    # ------------------------------------------------------------------
    def enable_expert_stats(self):
        """Start collecting per-token expert activation counts."""
        for layer in self.model.layers:
            if hasattr(layer, "feed_forward") and hasattr(layer.feed_forward, "gate"):
                layer.feed_forward.gate.collect_stats = True

    def disable_expert_stats(self):
        """Stop collecting expert activation counts."""
        for layer in self.model.layers:
            if hasattr(layer, "feed_forward") and hasattr(layer.feed_forward, "gate"):
                layer.feed_forward.gate.collect_stats = False

    def reset_expert_stats(self):
        """Clear all accumulated statistics."""
        for layer in self.model.layers:
            if hasattr(layer, "feed_forward") and hasattr(layer.feed_forward, "gate"):
                gate = layer.feed_forward.gate
                gate._stats_sum = 0.0
                gate._stats_sum_sq = 0.0
                gate._stats_count = 0

    def get_expert_activation_stats(self):
        """Return mean/std of experts activated per token, globally and per layer.

        Returns:
            dict with keys:
                - global: {mean_experts_per_token, std_experts_per_token, total_tokens_processed}
                - per_layer: {layer_0: {mean, std, count}, ...}
        """
        per_layer = {}
        global_sum = 0.0
        global_sum_sq = 0.0
        global_count = 0

        for layer_idx, layer in enumerate(self.model.layers):
            if hasattr(layer, "feed_forward") and hasattr(layer.feed_forward, "gate"):
                gate = layer.feed_forward.gate
                if gate._stats_count > 0:
                    mean = gate._stats_sum / gate._stats_count
                    mean_sq = gate._stats_sum_sq / gate._stats_count
                    # population std
                    std = max(0.0, mean_sq - mean ** 2) ** 0.5
                    per_layer[f"layer_{layer_idx}"] = {
                        "mean": round(mean, 4),
                        "std": round(std, 4),
                        "count": int(gate._stats_count),
                    }
                    global_sum += gate._stats_sum
                    global_sum_sq += gate._stats_sum_sq
                    global_count += gate._stats_count

        global_mean = global_sum / global_count if global_count > 0 else 0.0
        global_mean_sq = global_sum_sq / global_count if global_count > 0 else 0.0
        global_std = max(0.0, global_mean_sq - global_mean ** 2) ** 0.5 if global_count > 0 else 0.0

        return {
            "global": {
                "mean_experts_per_token": round(global_mean, 4),
                "std_experts_per_token": round(global_std, 4),
                "total_tokens_processed": int(global_count),
            },
            "per_layer": per_layer,
        }