"""PyTorch gLM2 model.

Minimal HuggingFace port of tattabio/gLM2 with three attention implementations
(eager, sdpa, flash_attention_2) and standard HF outputs.

Architecture is unchanged from the upstream `tattabio/gLM2_*` checkpoints
(RMSNorm, rotary position embeddings, fused QKV, SwiGLU MLP). Weight names
match upstream so the same `model.safetensors` loads cleanly.
"""

import math
from typing import Optional, Tuple, Union

import torch
import torch.nn.functional as F
from torch import nn
from torch.nn import CrossEntropyLoss

from transformers.modeling_outputs import BaseModelOutput, MaskedLMOutput
from transformers.modeling_utils import PreTrainedModel
from transformers.utils import logging

from .configuration_glm2 import gLM2Config

logger = logging.get_logger(__name__)


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_emb_torch(
    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
) -> torch.Tensor:
    """Apply rotary embeddings to `x`.

    Args:
        x: (batch, seqlen, nheads, headdim)
        cos, sin: (seqlen, rotary_dim / 2) - rotary_dim must equal headdim.
    """
    seqlen = x.shape[1]
    cos = cos[:seqlen]
    sin = sin[:seqlen]
    cos = cos.to(x.dtype)
    sin = sin.to(x.dtype)
    cos = cos.repeat_interleave(2, dim=-1) if False else torch.cat([cos, cos], dim=-1)
    sin = torch.cat([sin, sin], dim=-1)
    cos = cos.unsqueeze(-2)
    sin = sin.unsqueeze(-2)
    return x * cos + rotate_half(x) * sin


class RotaryEmbedding(nn.Module):
    """Rotary position embeddings.

    Identical numerics to the upstream `tattabio/gLM2_*` `RotaryEmbedding`
    (non-interleaved, base 10000, no scaling), simplified to the path actually
    used by the released checkpoints.
    """

    def __init__(self, dim: int, base: float = 10000.0):
        super().__init__()
        self.dim = dim
        self.base = float(base)
        inv_freq = 1.0 / (
            self.base
            ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)
        )
        self.register_buffer("inv_freq", inv_freq, persistent=False)

        self._seq_len_cached = 0
        self._cos_cached: Optional[torch.Tensor] = None
        self._sin_cached: Optional[torch.Tensor] = None

    def _update_cache(self, seqlen: int, device: torch.device, dtype: torch.dtype) -> None:
        if (
            seqlen > self._seq_len_cached
            or self._cos_cached is None
            or self._cos_cached.device != device
            or self._cos_cached.dtype != dtype
        ):
            self._seq_len_cached = seqlen
            inv_freq = self.inv_freq
            if inv_freq.dtype != torch.float32:
                inv_freq = 1.0 / (
                    self.base
                    ** (
                        torch.arange(0, self.dim, 2, device=device, dtype=torch.float32)
                        / self.dim
                    )
                )
            t = torch.arange(seqlen, device=device, dtype=torch.float32)
            freqs = torch.outer(t, inv_freq.to(device=device, dtype=torch.float32))
            self._cos_cached = torch.cos(freqs).to(dtype)
            self._sin_cached = torch.sin(freqs).to(dtype)

    def forward(self, qkv: torch.Tensor) -> torch.Tensor:
        """Apply rotary embeddings to q and k. v is left untouched.

        Args:
            qkv: (batch, seqlen, 3, nheads, headdim)
        """
        seqlen = qkv.shape[1]
        self._update_cache(seqlen, device=qkv.device, dtype=qkv.dtype)
        cos = self._cos_cached
        sin = self._sin_cached
        q_rot = apply_rotary_emb_torch(qkv[:, :, 0], cos, sin)
        k_rot = apply_rotary_emb_torch(qkv[:, :, 1], cos, sin)
        return torch.stack((q_rot, k_rot, qkv[:, :, 2]), dim=2)


def rmsnorm_func(
    hidden_states: torch.Tensor, weight: torch.Tensor, variance_epsilon: torch.Tensor
) -> torch.Tensor:
    input_dtype = hidden_states.dtype
    hidden_states = hidden_states.to(torch.float32)
    variance = hidden_states.pow(2).mean(-1, keepdim=True)
    hidden_states = hidden_states * torch.rsqrt(variance + variance_epsilon)
    return (weight * hidden_states).to(input_dtype)


class RMSNorm(nn.Module):
    """Root-mean-square layer norm."""

    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.register_buffer(
            "variance_epsilon", torch.tensor(eps), persistent=False
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return rmsnorm_func(hidden_states, self.weight, self.variance_epsilon)


class gLM2Attention(nn.Module):
    """Eager multi-head attention with rotary embeddings."""

    def __init__(self, config: gLM2Config):
        super().__init__()
        self.n_heads = config.heads
        self.head_dim = config.dim // config.heads
        self.dim = config.dim

        self.wqkv = nn.Linear(config.dim, self.n_heads * self.head_dim * 3, bias=False)
        self.wo = nn.Linear(self.n_heads * self.head_dim, config.dim, bias=False)
        self.rotary_emb = RotaryEmbedding(self.head_dim)

    def _qkv(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        bsz, seqlen, _ = x.shape
        qkv = self.wqkv(x).view(bsz, seqlen, 3, self.n_heads, self.head_dim)
        qkv = self.rotary_emb(qkv)
        # qkv: (B, S, 3, H, D) -> (B, H, S, D) for q,k,v
        qkv = qkv.permute(0, 3, 2, 1, 4)  # (B, H, 3, S, D)
        q = qkv[:, :, 0]
        k = qkv[:, :, 1]
        v = qkv[:, :, 2]
        return q, k, v

    def _output(self, attn_out: torch.Tensor) -> torch.Tensor:
        # attn_out: (B, H, S, D) -> (B, S, H*D)
        bsz, _, seqlen, _ = attn_out.shape
        out = attn_out.permute(0, 2, 1, 3).contiguous().view(bsz, seqlen, self.n_heads * self.head_dim)
        return self.wo(out)

    def forward(
        self,
        x: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        q, k, v = self._qkv(x)
        scale = 1.0 / math.sqrt(self.head_dim)
        # (B, H, S, S). Compute in fp32 for numerical stability under bf16/fp16
        # (matches what flash-attn / SDPA do internally).
        scores = torch.matmul(q.float(), k.float().transpose(-2, -1)) * scale
        if attention_mask is not None:
            mask = attention_mask[:, None, None, :]
            scores = scores.masked_fill(mask == 0, torch.finfo(scores.dtype).min)
        attn = scores.softmax(dim=-1)
        attn_for_return = attn.to(q.dtype) if output_attentions else None
        context = torch.matmul(attn, v.float()).to(q.dtype)
        return self._output(context), attn_for_return


class gLM2SdpaAttention(gLM2Attention):
    """SDPA-backed attention. Falls back to eager when output_attentions=True."""

    def forward(
        self,
        x: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        if output_attentions:
            return super().forward(x, attention_mask=attention_mask, output_attentions=True)

        q, k, v = self._qkv(x)
        attn_mask = None
        if attention_mask is not None:
            # SDPA wants (B, 1, 1, S) bool mask where True = attend.
            attn_mask = attention_mask[:, None, None, :].bool()
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
        return self._output(out), None


class gLM2FlashAttention2(gLM2Attention):
    """flash-attn 2 backed attention. Falls back to eager when output_attentions=True."""

    def forward(
        self,
        x: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        if output_attentions:
            return super().forward(x, attention_mask=attention_mask, output_attentions=True)

        try:
            from flash_attn import flash_attn_func, flash_attn_varlen_func
            from flash_attn.bert_padding import pad_input, unpad_input
        except ImportError as e:
            raise ImportError(
                "flash_attn is required for attn_implementation='flash_attention_2'. "
                "Install with: pip install flash-attn --no-build-isolation"
            ) from e

        bsz, seqlen, _ = x.shape
        qkv = self.wqkv(x).view(bsz, seqlen, 3, self.n_heads, self.head_dim)
        qkv = self.rotary_emb(qkv)
        # flash-attn wants (B, S, H, D) per q/k/v.
        q = qkv[:, :, 0]
        k = qkv[:, :, 1]
        v = qkv[:, :, 2]

        orig_dtype = q.dtype
        if q.dtype not in (torch.float16, torch.bfloat16):
            q = q.to(torch.bfloat16)
            k = k.to(torch.bfloat16)
            v = v.to(torch.bfloat16)

        if attention_mask is not None and (attention_mask == 0).any():
            attention_mask_bool = attention_mask.bool()  # True = attend
            q_unpad, indices_q, cu_q, max_q, _ = unpad_input(q, attention_mask_bool)
            k_unpad, _, cu_k, max_k, _ = unpad_input(k, attention_mask_bool)
            v_unpad, _, _, _, _ = unpad_input(v, attention_mask_bool)
            out_unpad = flash_attn_varlen_func(
                q_unpad,
                k_unpad,
                v_unpad,
                cu_seqlens_q=cu_q,
                cu_seqlens_k=cu_k,
                max_seqlen_q=max_q,
                max_seqlen_k=max_k,
                causal=False,
            )
            out = pad_input(out_unpad, indices_q, bsz, seqlen)
        else:
            out = flash_attn_func(q, k, v, causal=False)

        out = out.to(orig_dtype)
        out = out.contiguous().view(bsz, seqlen, self.n_heads * self.head_dim)
        return self.wo(out), None


GLM2_ATTENTION_CLASSES = {
    "eager": gLM2Attention,
    "sdpa": gLM2SdpaAttention,
    "flash_attention_2": gLM2FlashAttention2,
}


class FeedForward(nn.Module):
    """SwiGLU MLP."""

    def __init__(
        self,
        dim: int,
        hidden_dim: int,
        multiple_of: int,
        ffn_dim_multiplier: Optional[float],
    ):
        super().__init__()
        hidden_dim = int(2 * hidden_dim / 3)
        if ffn_dim_multiplier is not None:
            hidden_dim = int(ffn_dim_multiplier * hidden_dim)
        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)

        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
        self.w3 = nn.Linear(dim, hidden_dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


class TransformerBlock(nn.Module):
    """Pre-norm transformer block."""

    def __init__(self, config: gLM2Config):
        super().__init__()
        attn_impl = getattr(config, "_attn_implementation", "eager")
        attn_cls = GLM2_ATTENTION_CLASSES[attn_impl]
        self.attention = attn_cls(config)
        self.feed_forward = FeedForward(
            dim=config.dim,
            hidden_dim=4 * config.dim,
            multiple_of=config.swiglu_multiple_of,
            ffn_dim_multiplier=config.ffn_dim_multiplier,
        )
        self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps)
        self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)

    def forward(
        self,
        x: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        attn_out, attn_weights = self.attention(
            self.attention_norm(x),
            attention_mask=attention_mask,
            output_attentions=output_attentions,
        )
        h = x + attn_out
        out = h + self.feed_forward(self.ffn_norm(h))
        return out, attn_weights


class TransformerLayers(nn.Module):
    def __init__(self, config: gLM2Config):
        super().__init__()
        self.config = config
        self.layers = nn.ModuleList(
            [TransformerBlock(config) for _ in range(config.depth)]
        )

    def forward(
        self,
        x: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        output_hidden_states: bool = False,
        output_attentions: bool = False,
    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, ...]], Optional[Tuple[torch.Tensor, ...]]]:
        if x.shape[-1] != self.config.dim:
            raise ValueError(
                f"Input feature dim should be {self.config.dim}, but input has shape {x.shape}"
            )
        all_hidden_states: list = []
        all_attentions: list = []
        if output_hidden_states:
            all_hidden_states.append(x)
        for layer in self.layers:
            x, attn_weights = layer(
                x, attention_mask=attention_mask, output_attentions=output_attentions
            )
            if output_hidden_states:
                all_hidden_states.append(x)
            if output_attentions:
                all_attentions.append(attn_weights)
        hidden_tuple = tuple(all_hidden_states) if output_hidden_states else None
        attn_tuple = tuple(all_attentions) if output_attentions else None
        return x, hidden_tuple, attn_tuple


class gLM2PreTrainedModel(PreTrainedModel):
    """Base class for gLM2 weight init / from_pretrained dispatch."""

    config_class = gLM2Config
    base_model_prefix = "glm2"
    supports_gradient_checkpointing = False
    _supports_sdpa = True
    _supports_flash_attn_2 = True

    def _init_weights(self, module):
        std = getattr(self.config, "initializer_range", 0.02)
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, std=std)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, std=std)
            if module.padding_idx is not None:
                with torch.no_grad():
                    module.weight[module.padding_idx].zero_()
        elif isinstance(module, RotaryEmbedding):
            inv_freq = 1.0 / (
                module.base
                ** (
                    torch.arange(
                        0, module.dim, 2, device=module.inv_freq.device, dtype=torch.float32
                    )
                    / module.dim
                )
            )
            with torch.no_grad():
                module.inv_freq.copy_(inv_freq)
        elif isinstance(module, RMSNorm):
            with torch.no_grad():
                module.variance_epsilon.fill_(self.config.norm_eps)


class gLM2Model(gLM2PreTrainedModel):
    """gLM2 backbone (token embedding + transformer encoder)."""

    def __init__(self, config: gLM2Config):
        super().__init__(config)
        self.config = config
        self.tok_embeddings = nn.Embedding(config.vocab_size, config.dim)
        self.encoder = TransformerLayers(config)
        self.post_init()

    def get_input_embeddings(self) -> nn.Embedding:
        return self.tok_embeddings

    def set_input_embeddings(self, value: nn.Embedding) -> None:
        self.tok_embeddings = value

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor, ...], BaseModelOutput]:
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        output_attentions = (
            output_attentions
            if output_attentions is not None
            else self.config.output_attentions
        )
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        h = self.tok_embeddings(input_ids)
        sequence_output, all_hidden_states, all_attentions = self.encoder(
            h,
            attention_mask=attention_mask,
            output_hidden_states=bool(output_hidden_states),
            output_attentions=bool(output_attentions),
        )

        if not return_dict:
            return tuple(
                v
                for v in (sequence_output, all_hidden_states, all_attentions)
                if v is not None
            )

        return BaseModelOutput(
            last_hidden_state=sequence_output,
            hidden_states=all_hidden_states,
            attentions=all_attentions,
        )


class gLM2LMHead(nn.Module):
    def __init__(self, config: gLM2Config):
        super().__init__()
        self.norm = RMSNorm(config.dim, eps=config.norm_eps)
        self.proj_output = nn.Linear(config.dim, config.vocab_size, bias=False)

    def forward(self, features: torch.Tensor) -> torch.Tensor:
        return self.proj_output(self.norm(features))


class gLM2ForMaskedLM(gLM2PreTrainedModel):
    """gLM2 with the masked-language-modeling head."""

    _tied_weights_keys = []

    def __init__(self, config: gLM2Config):
        super().__init__(config)
        self.glm2 = gLM2Model(config)
        self.lm_head = gLM2LMHead(config)
        self.post_init()

    def get_output_embeddings(self) -> nn.Linear:
        return self.lm_head.proj_output

    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
        self.lm_head.proj_output = new_embeddings

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor, ...], MaskedLMOutput]:
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        outputs = self.glm2(
            input_ids,
            attention_mask=attention_mask,
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
            return_dict=True,
        )
        sequence_output = outputs.last_hidden_state
        prediction_scores = self.lm_head(sequence_output)

        masked_lm_loss = None
        if labels is not None:
            loss_fct = CrossEntropyLoss()
            labels = labels.to(prediction_scores.device)
            masked_lm_loss = loss_fct(
                prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)
            )

        if not return_dict:
            output = (prediction_scores,)
            if outputs.hidden_states is not None:
                output = output + (outputs.hidden_states,)
            if outputs.attentions is not None:
                output = output + (outputs.attentions,)
            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output

        return MaskedLMOutput(
            loss=masked_lm_loss,
            logits=prediction_scores,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )