"""Tokle: decoder-only transformer with RMSNorm, RoPE, GQA and SwiGLU.

Module and parameter names follow the usual Transformers layout.

Works on Transformers 4.x and 5.x; the v5 differences are marked inline.
"""
import math
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import __version__ as _transformers_version
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import BaseModelOutput, CausalLMOutput
from transformers.modeling_utils import PreTrainedModel

try:  # as a Transformers dynamic module (trust_remote_code)
    from .configuration_tokle import TokleConfig
except ImportError:  # running the files straight out of a checkout
    from configuration_tokle import TokleConfig


# ----------------------------------------------------------------------------- norms and rope
class TokleRMSNorm(nn.Module):
    def __init__(self, d, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d))

    def forward(self, x):
        # Always reduce in fp32: in fp16 the mean of squares underflows on the
        # long tail and the norm comes back slightly wrong.
        dt = x.dtype
        x = x.float()
        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return (x * self.weight.float()).to(dt)

    def extra_repr(self):
        return f"{tuple(self.weight.shape)}, eps={self.eps}"


def build_rope(head_dim, max_seq_len, theta, device=None):
    inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32,
                                        device=device) / head_dim))
    t = torch.arange(max_seq_len, dtype=torch.float32, device=device)
    f = torch.outer(t, inv)
    return torch.cos(f), torch.sin(f)


def apply_rope(x, cos, sin):
    """Half-split rotation (first half against second), not the interleaved
    variant. The two are not interchangeable -- swapping them degrades output
    without raising anything."""
    b, h, t, d = x.shape
    x1, x2 = x.float().chunk(2, dim=-1)
    c = cos[:t].view(1, 1, t, d // 2).to(x1.device)
    s = sin[:t].view(1, 1, t, d // 2).to(x1.device)
    return torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], dim=-1).to(x.dtype)


# ----------------------------------------------------------------------------- blocks
class TokleAttention(nn.Module):
    def __init__(self, config: TokleConfig, layer_idx: int):
        super().__init__()
        self.layer_idx = layer_idx
        self.n_head = config.n_head
        self.n_kv_head = config.n_kv_head
        self.head_dim = config.head_dim
        self.rep = config.n_head // config.n_kv_head
        self.q_proj = nn.Linear(config.d_model, config.n_head * config.head_dim, bias=False)
        self.k_proj = nn.Linear(config.d_model, config.n_kv_head * config.head_dim, bias=False)
        self.v_proj = nn.Linear(config.d_model, config.n_kv_head * config.head_dim, bias=False)
        self.o_proj = nn.Linear(config.n_head * config.head_dim, config.d_model, bias=False)
        self.q_norm = TokleRMSNorm(config.head_dim, config.norm_eps)
        self.k_norm = TokleRMSNorm(config.head_dim, config.norm_eps)

    def forward(self, x, cos, sin, attn_mask):
        b, t, _ = x.shape
        q = self.q_proj(x).view(b, t, self.n_head, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(b, t, self.n_kv_head, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(b, t, self.n_kv_head, self.head_dim).transpose(1, 2)
        # QK-norm is applied before rope.
        q = apply_rope(self.q_norm(q), cos, sin)
        k = apply_rope(self.k_norm(k), cos, sin)
        k = k.repeat_interleave(self.rep, dim=1)
        v = v.repeat_interleave(self.rep, dim=1)
        if attn_mask is None:
            # Nothing padded: let SDPA build the causal mask itself, which is
            # the fast kernel path and what an unpadded forward always hits.
            o = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        else:
            o = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask.to(q.dtype))
        return self.o_proj(o.transpose(1, 2).reshape(b, t, -1))


class TokleMLP(nn.Module):
    def __init__(self, config: TokleConfig):
        super().__init__()
        self.gate_proj = nn.Linear(config.d_model, config.ffn_hidden, bias=False)
        self.up_proj = nn.Linear(config.d_model, config.ffn_hidden, bias=False)
        self.down_proj = nn.Linear(config.ffn_hidden, config.d_model, bias=False)

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


class TokleDecoderLayer(nn.Module):
    def __init__(self, config: TokleConfig, layer_idx: int):
        super().__init__()
        self.input_layernorm = TokleRMSNorm(config.d_model, config.norm_eps)
        self.self_attn = TokleAttention(config, layer_idx)
        self.post_attention_layernorm = TokleRMSNorm(config.d_model, config.norm_eps)
        self.mlp = TokleMLP(config)

    def forward(self, x, cos, sin, attn_mask):
        x = x + self.self_attn(self.input_layernorm(x), cos, sin, attn_mask)
        return x + self.mlp(self.post_attention_layernorm(x))


# ----------------------------------------------------------------------------- model
class ToklePreTrainedModel(PreTrainedModel):
    config_class = TokleConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["TokleDecoderLayer"]
    _supports_sdpa = True

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, 0.0, 0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, 0.0, 0.02)
        elif isinstance(module, TokleRMSNorm):
            nn.init.ones_(module.weight)
        # Residual-output projections get scaled down by depth so the residual
        # stream does not blow up at init (GPT-2 trick).
        if isinstance(module, (TokleAttention, TokleMLP)):
            out = module.o_proj if isinstance(module, TokleAttention) else module.down_proj
            nn.init.normal_(out.weight, 0.0, 0.02 / math.sqrt(2 * self.config.n_layer))


class TokleModel(ToklePreTrainedModel):
    def __init__(self, config: TokleConfig):
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model)
        self.layers = nn.ModuleList(
            [TokleDecoderLayer(config, i) for i in range(config.n_layer)])
        self.norm = TokleRMSNorm(config.d_model, config.norm_eps)
        # A plain dict, not a buffer. Transformers 5 builds models on the meta
        # device and fills only tensors the weights file contains, so a
        # non-persistent buffer would come back as uninitialised memory.
        # Keyed by device, so .to()/.cuda() need not move anything.
        self._rope_cache = {}
        self.gradient_checkpointing = False
        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

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

    def rope(self, device):
        key = str(device)
        if key not in self._rope_cache:
            self._rope_cache[key] = build_rope(
                self.config.head_dim, self.config.max_seq_len,
                self.config.rope_theta, device=device)
        return self._rope_cache[key]

    def _causal_pad_mask(self, attention_mask, t, dtype, device):
        """Additive (b, 1, 1, t) padding mask folded together with the causal
        mask, or None when nothing is padded so the caller can take the
        is_causal fast path."""
        if attention_mask is None:
            return None
        attention_mask = attention_mask.to(torch.bool)
        if bool(attention_mask.all()):
            return None
        m = torch.zeros(attention_mask.shape, dtype=dtype, device=device)
        m = m.masked_fill(~attention_mask, torch.finfo(dtype).min)[:, None, None, :]
        # Under causal attention with right padding, position 0 always sees
        # itself, so no row ends up fully masked (which would give NaNs).
        causal = torch.ones(t, t, dtype=torch.bool, device=device).triu(1)
        return m.masked_fill(causal, torch.finfo(dtype).min)

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, BaseModelOutput]:
        # getattr, not config.use_return_dict: that property is deprecated in
        # Transformers 5 and warns on every call.
        return_dict = (return_dict if return_dict is not None
                       else getattr(self.config, "return_dict", True))
        output_hidden_states = (output_hidden_states if output_hidden_states is not None
                                else self.config.output_hidden_states)
        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("Pass exactly one of input_ids or inputs_embeds.")

        x = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds
        t = x.shape[1]
        if t > self.config.max_seq_len:
            raise ValueError(
                f"Sequence length {t} exceeds max_seq_len {self.config.max_seq_len}.")

        cos, sin = self.rope(x.device)
        mask = self._causal_pad_mask(attention_mask, t, x.dtype, x.device)

        hidden_states = () if output_hidden_states else None
        for layer in self.layers:
            if output_hidden_states:
                hidden_states += (x,)
            if self.gradient_checkpointing and self.training:
                x = self._gradient_checkpointing_func(layer.__call__, x, cos, sin, mask)
            else:
                x = layer(x, cos, sin, mask)
        x = self.norm(x)
        if output_hidden_states:
            hidden_states += (x,)

        if not return_dict:
            return tuple(v for v in (x, hidden_states) if v is not None)
        return BaseModelOutput(last_hidden_state=x, hidden_states=hidden_states)


# Transformers 5 turned _tied_weights_keys from a list of names into a
# {tied: source} mapping and calls .keys() on it, so the old form raises.
_TRANSFORMERS_V5 = int(_transformers_version.split(".")[0]) >= 5


class TokleForCausalLM(ToklePreTrainedModel, GenerationMixin):
    _tied_weights_keys = ({"lm_head.weight": "model.embed_tokens.weight"}
                          if _TRANSFORMERS_V5 else ["lm_head.weight"])

    def __init__(self, config: TokleConfig):
        super().__init__(config)
        self.model = TokleModel(config)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.post_init()

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

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

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def get_decoder(self):
        return self.model

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, CausalLMOutput]:
        return_dict = (return_dict if return_dict is not None
                       else getattr(self.config, "return_dict", True))
        out = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            output_hidden_states=output_hidden_states,
            return_dict=True,
        )
        logits = self.lm_head(out.last_hidden_state)
        # Rows past real_vocab_size only pad the matrix to a multiple of 64 and
        # carry no meaning, so keep them out of sampling and out of any
        # log-likelihood an evaluation harness sums.
        rv = self.config.real_vocab_size
        if self.config.mask_padded_vocab_logits and rv is not None and rv < self.config.vocab_size:
            logits = logits.clone()
            logits[..., rv:] = float("-inf")

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

        if not return_dict:
            return (loss, logits) if loss is not None else (logits,)
        return CausalLMOutput(loss=loss, logits=logits, hidden_states=out.hidden_states)