import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers.modeling_utils import PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import MoeCausalLMOutputWithPast
from transformers.cache_utils import Cache, DynamicCache

from .configuration_dynamicmind_moe import DynamicMindMoEConfig


class DynamicMindRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, x):
        dtype = x.dtype
        x = x.float()
        x = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return (self.weight * x).to(dtype)


class DynamicMindRotaryEmbedding(nn.Module):
    """RoPE with a cached inv_freq.

    The dense model rebuilt inv_freq on every forward of every layer; caching it
    removes 9 redundant allocations per step.

    inv_freq is a constant derived from config, so persistent=False looks
    correct — but from_pretrained materialises tensors straight from the
    checkpoint onto meta-device modules, never running __init__'s value nor
    _load_from_state_dict for it. A non-persistent buffer therefore survives
    loading as uninitialised `torch.empty` garbage, silently scrambling RoPE:
    measured 833 vs 908 Elo on identical weights. Persisting the 16 floats is
    the only variant that loads correctly through every path.
    """

    def __init__(self, head_dim, rope_theta, max_position_embeddings):
        super().__init__()
        self.head_dim = head_dim
        self.rope_theta = rope_theta
        self.register_buffer("inv_freq", self._compute(), persistent=True)
        self.max_seq_len_cached = 0

    def _compute(self):
        return 1.0 / (self.rope_theta ** (
            torch.arange(0, self.head_dim, 2).float() / self.head_dim))

    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
        with torch.no_grad():
            self.inv_freq.copy_(self._compute().to(self.inv_freq.device))

    def forward(self, x, position_ids):
        freqs = position_ids[:, :, None].float() * self.inv_freq[None, None, :]
        return freqs.cos().to(x.dtype), freqs.sin().to(x.dtype)


def apply_rope(q, k, cos, sin):
    cos = cos[:, None, :, :]
    sin = sin[:, None, :, :]

    def rotate(x):
        even, odd = x[..., 0::2], x[..., 1::2]
        return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)

    return rotate(q), rotate(k)


class DynamicMindAttention(nn.Module):
    def __init__(self, config, layer_idx):
        super().__init__()
        self.layer_idx = layer_idx
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.hidden_size // config.num_attention_heads
        self.attention_dropout = config.attention_dropout

        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)

    def forward(self, x, cos, sin, attention_mask=None, past_key_values=None, cache_position=None):
        bsz, q_len, _ = x.shape

        q = self.q_proj(x).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)

        q, k = apply_rope(q, k, cos, sin)

        if past_key_values is not None:
            k, v = past_key_values.update(k, v, self.layer_idx, {"cache_position": cache_position})

        if self.num_kv_heads != self.num_heads:
            repeats = self.num_heads // self.num_kv_heads
            k = k.repeat_interleave(repeats, dim=1)
            v = v.repeat_interleave(repeats, dim=1)

        is_causal = attention_mask is None and q_len > 1
        y = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=attention_mask,
            dropout_p=self.attention_dropout if self.training else 0.0,
            is_causal=is_causal,
        )
        y = y.transpose(1, 2).contiguous().view(bsz, q_len, -1)
        return self.o_proj(y)


class DynamicMindMLP(nn.Module):
    def __init__(self, hidden_size, intermediate_size):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class DynamicMindMoE(nn.Module):
    """Shared expert + top-k routed fine-grained experts.

    The shared expert runs on every token and absorbs knowledge common to all
    inputs, so the routed experts are free to specialise instead of each
    re-learning the same basics.
    """

    def __init__(self, config):
        super().__init__()
        self.num_routed = config.num_routed_experts
        self.top_k = config.num_experts_per_token
        self.norm_topk_prob = config.norm_topk_prob
        self.aux_free = config.use_aux_loss_free_balancing

        self.experts = nn.ModuleList([
            DynamicMindMLP(config.hidden_size, config.moe_intermediate_size)
            for _ in range(self.num_routed)
        ])
        self.shared_experts = nn.ModuleList([
            DynamicMindMLP(config.hidden_size, config.moe_intermediate_size)
            for _ in range(config.num_shared_experts)
        ])

        self.router = nn.Linear(config.hidden_size, self.num_routed, bias=False)

        # Aux-loss-free balancing: a per-expert bias nudged toward even load.
        # It steers selection only — never the combining weights — so it costs
        # no gradient interference, unlike an auxiliary loss.
        self.register_buffer("expert_bias", torch.zeros(self.num_routed), persistent=True)
        self.bias_update_rate = config.router_bias_update_rate

    def forward(self, x):
        bsz, seq_len, hidden = x.shape
        flat = x.view(-1, hidden)
        n_tokens = flat.size(0)

        logits = self.router(flat)                       # [T, E]
        probs = F.softmax(logits, dim=-1, dtype=torch.float)

        scores = probs + self.expert_bias if self.aux_free else probs
        _, topk_idx = torch.topk(scores, self.top_k, dim=-1)
        topk_w = probs.gather(-1, topk_idx)              # weights from unbiased probs
        if self.norm_topk_prob:
            topk_w = topk_w / topk_w.sum(dim=-1, keepdim=True).clamp_min(1e-9)
        topk_w = topk_w.to(x.dtype)

        out = torch.zeros_like(flat)
        for expert in self.shared_experts:
            out = out + expert(flat)

        # one-hot over experts -> per-expert token lists
        mask = torch.zeros(n_tokens, self.num_routed, dtype=torch.bool, device=x.device)
        mask.scatter_(1, topk_idx, True)
        load = mask.sum(0)

        for e in range(self.num_routed):
            idx = mask[:, e].nonzero(as_tuple=True)[0]
            if idx.numel() == 0:
                continue
            slot = (topk_idx[idx] == e).float().argmax(dim=-1)
            w = topk_w[idx].gather(-1, slot[:, None])
            out.index_add_(0, idx, self.experts[e](flat[idx]) * w)

        if self.training and self.aux_free:
            with torch.no_grad():
                target = n_tokens * self.top_k / self.num_routed
                self.expert_bias += self.bias_update_rate * (target - load.float()).sign()

        # Reported for logging even when aux-free balancing is on.
        frac_tokens = load.float() / (n_tokens * self.top_k)
        frac_probs = probs.mean(dim=0)
        aux_loss = self.num_routed * (frac_tokens * frac_probs).sum()
        z_loss = torch.logsumexp(logits.float(), dim=-1).pow(2).mean()

        return out.view(bsz, seq_len, hidden), aux_loss, z_loss, load


class DynamicMindBlock(nn.Module):
    def __init__(self, config, layer_idx):
        super().__init__()
        self.input_layernorm = DynamicMindRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.self_attn = DynamicMindAttention(config, layer_idx)
        self.post_attention_layernorm = DynamicMindRMSNorm(config.hidden_size, config.rms_norm_eps)

        self.is_moe = layer_idx >= config.first_k_dense_layers
        if self.is_moe:
            self.mlp = DynamicMindMoE(config)
        else:
            self.mlp = DynamicMindMLP(config.hidden_size, config.intermediate_size)

    def forward(self, x, cos, sin, attention_mask=None, past_key_values=None, cache_position=None):
        x = x + self.self_attn(self.input_layernorm(x), cos, sin,
                               attention_mask, past_key_values, cache_position)
        h = self.post_attention_layernorm(x)
        if self.is_moe:
            delta, aux, z, load = self.mlp(h)
            return x + delta, aux, z, load
        return x + self.mlp(h), None, None, None


class DynamicMindMoEPreTrainedModel(PreTrainedModel):
    config_class = DynamicMindMoEConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["DynamicMindBlock"]

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)


class DynamicMindMoEForCausalLM(DynamicMindMoEPreTrainedModel, GenerationMixin):
    _tied_weights_keys = {"lm_head.weight": "embed_tokens.weight"}

    def __init__(self, config):
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList([
            DynamicMindBlock(config, i) for i in range(config.num_hidden_layers)
        ])
        self.norm = DynamicMindRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.rotary = DynamicMindRotaryEmbedding(
            config.hidden_size // config.num_attention_heads,
            config.rope_theta,
            config.max_position_embeddings,
        )
        if config.tie_word_embeddings:
            self.lm_head.weight = self.embed_tokens.weight
        self.post_init()

    def tie_weights(self, *args, **kwargs):
        if getattr(self.config, "tie_word_embeddings", True):
            self.lm_head.weight = self.embed_tokens.weight

    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=None, labels=None, use_cache=True, **kwargs):
        # cache_position is read from kwargs rather than declared: transformers
        # warns about remote-code models whose signature expects it, and plans
        # to stop passing it. It is derived below whenever it is absent.
        cache_position = kwargs.get("cache_position")
        x = self.embed_tokens(input_ids)

        if use_cache and past_key_values is None:
            past_key_values = DynamicCache()
        past_len = past_key_values.get_seq_length() if isinstance(past_key_values, Cache) else 0

        if cache_position is None:
            cache_position = torch.arange(past_len, past_len + x.size(1), device=x.device)
        if position_ids is None:
            position_ids = cache_position[None, :]

        cos, sin = self.rotary(x, position_ids)

        causal_mask = None
        if x.size(1) > 1:
            total = past_len + x.size(1)
            causal = torch.tril(torch.ones(x.size(1), total, dtype=torch.bool, device=x.device),
                                diagonal=past_len)
            causal_mask = torch.zeros(x.size(1), total, dtype=x.dtype, device=x.device)
            causal_mask.masked_fill_(~causal, torch.finfo(x.dtype).min)
            causal_mask = causal_mask[None, None, :, :]

        aux_total = x.new_zeros(())
        z_total = x.new_zeros(())
        loads = []
        for layer in self.layers:
            x, aux, z, load = layer(x, cos, sin, causal_mask, past_key_values, cache_position)
            if aux is not None:
                aux_total = aux_total + aux
                z_total = z_total + z
                loads.append(load)

        logits = self.lm_head(self.norm(x))

        loss = None
        if labels is not None:
            shift_labels = torch.cat(
                [labels[:, 1:], labels.new_full((labels.size(0), 1), -100)], dim=1
            )
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), shift_labels.view(-1))
            n_moe = max(len(loads), 1)
            loss = loss + self.config.router_aux_loss_coef * aux_total / n_moe
            loss = loss + self.config.router_z_loss_coef * z_total / n_moe

        return MoeCausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=past_key_values if use_cache else None,
            aux_loss=aux_total / max(len(loads), 1) if loads else None,
        )

    def expert_load(self):
        """Per-layer expert token counts from the last forward, for monitoring."""
        return [m.expert_bias for m in self.modules() if isinstance(m, DynamicMindMoE)]

    def state_dict(self, *args, **kwargs):
        sd = super().state_dict(*args, **kwargs)
        if getattr(self.config, "tie_word_embeddings", True):
            for k in list(sd.keys()):
                if k == "lm_head.weight" or k.endswith(".lm_head.weight"):
                    del sd[k]
        return sd