"""
modeling_ivme.py — standalone HuggingFace-compatible model definition for
IvmeLabs Ivme-Conversate-XL-v1.

This file is self-contained (no dependency on train_xl_v1.py or liger-kernel)
so it works as a `trust_remote_code=True` HF Hub model: anyone can
`AutoModelForCausalLM.from_pretrained("IvmeLabs/Ivme-Conversate-XL-v1",
trust_remote_code=True)` and get a working model with no extra setup.

Architecture: dense decoder-only transformer, 125.6M params.
  12 layers, hidden=768, 12 heads, head_dim=64, SwiGLU ffn=3072, vocab=16000,
  RoPE (theta=10000), RMSNorm pre-norm, tied embeddings, no bias.

This is a plain-PyTorch port of the architecture trained in train_xl_v1.py.
Liger-Kernel fused ops (RMSNorm/RoPE/CrossEntropy) were used during training
for speed on ROCm, but are intentionally NOT a dependency here — inference
correctness only requires the same math, not the same fused kernels, and a
published model repo should load with just `torch` + `transformers`
installed, without requiring users to also have liger-kernel (and by
extension a matching ROCm/Triton setup) just to run inference.

Attention weights: by default, attention uses F.scaled_dot_product_attention
(SDPA), a fused kernel that never materializes the [B, H, T, T] attention
probability matrix — so output_attentions=True cannot return anything on
that path. When output_attentions=True is requested, this file falls back
to a manual (non-fused, eager) attention implementation per layer that
does compute and return real softmax attention weights, at the cost of
speed/memory versus SDPA. This mirrors the fused-vs-eager tradeoff noted
in train_xl_v1.py for the Liger kernels used during training.
"""

import math
from typing import Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PretrainedConfig, PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast


class IvmeXLConfig(PretrainedConfig):
    model_type = "ivme_xl"

    def __init__(
        self,
        vocab_size: int = 16000,
        hidden_size: int = 768,
        n_layers: int = 12,
        n_heads: int = 12,
        head_dim: int = 64,
        ffn_dim: int = 3072,
        max_position_embeddings: int = 1024,
        rope_theta: float = 10000.0,
        norm_eps: float = 1e-5,
        tie_word_embeddings: bool = True,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.head_dim = head_dim
        self.ffn_dim = ffn_dim
        self.max_position_embeddings = max_position_embeddings
        self.rope_theta = rope_theta
        self.norm_eps = norm_eps
        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

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


def precompute_rope_freqs(head_dim: int, max_pos: int, theta: float,
                           device, dtype):
    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    t = torch.arange(max_pos, device=device).float()
    freqs = torch.outer(t, inv_freq)
    emb = torch.cat((freqs, freqs), dim=-1)
    return emb.cos().to(dtype), emb.sin().to(dtype)


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


def apply_rope(q, k, cos, sin):
    # q, k: [B, H, T, D]; cos/sin: [B, T, D] (batch dim included — see
    # train_xl_v1.py's Attention.forward comment for why this matters: a
    # missing batch dim here caused a real GPU memory-fault bug during
    # training with the fused Liger RoPE kernel. Kept as [B, T, D] here too
    # for consistency with the trained checkpoint's expected convention,
    # even though the plain-PyTorch path is more forgiving of the 2D case.
    cos = cos.unsqueeze(1)
    sin = sin.unsqueeze(1)
    q_out = (q * cos) + (rotate_half(q) * sin)
    k_out = (k * cos) + (rotate_half(k) * sin)
    return q_out, k_out


def _expand_padding_mask(
    attention_mask: torch.Tensor, seq_len: int, dtype: torch.dtype, mask_value: float
) -> Optional[torch.Tensor]:
    """Convert a raw HF-style padding mask into an additive bias broadcastable
    against [B, H, T, T] attention scores.

    HF convention: attention_mask has shape [B, T] with 1 = keep, 0 = pad.
    If it's already a >=3D additive-style mask (e.g. a caller passing a
    precomputed bias), pass it through unchanged rather than reinterpreting it.
    """
    if attention_mask is None:
        return None
    if attention_mask.dim() >= 3:
        # Already shaped as an additive bias (e.g. [B, 1, T, T] or [B, 1, 1, T]).
        return attention_mask.to(dtype)
    if attention_mask.dim() != 2 or attention_mask.shape[-1] != seq_len:
        # Shape doesn't match what we expect for a [B, T] padding mask;
        # don't guess-broadcast something that could silently misalign.
        return None
    # [B, T] (1=keep, 0=pad) -> [B, 1, 1, T] additive bias, applied to the
    # key dimension so padded *keys* get ~0 attention probability.
    bias = (1.0 - attention_mask.to(dtype)) * mask_value
    return bias[:, None, None, :]


class Attention(nn.Module):
    def __init__(self, config: IvmeXLConfig):
        super().__init__()
        self.n_heads = config.n_heads
        self.head_dim = config.head_dim
        self.hidden_size = config.hidden_size

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

    def forward(
        self,
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        B, T, C = x.shape
        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)

        cos_t = cos[:T].unsqueeze(0).expand(B, -1, -1)
        sin_t = sin[:T].unsqueeze(0).expand(B, -1, -1)
        q, k = apply_rope(q, k, cos_t, sin_t)

        attn_weights = None

        # Use a large-but-bounded finite negative number instead of -inf for
        # masking. -inf is dangerous here for two reasons: (1) if a row of
        # attn_scores ends up entirely masked, softmax over an all -inf row
        # is 0/0 -> NaN; (2) the causal mask and the padding mask can be
        # summed on the same position, and float32's true min
        # (~-3.4e38) would overflow to -inf when added to itself or to a
        # nonzero score, reintroducing the same NaN. -1e9 is comfortably
        # below any realistic logit magnitude yet stays finite under
        # addition, so softmax reliably drives masked positions to ~0
        # probability without ever hitting NaN/inf.
        mask_value = -1e9

        if output_attentions:
            # Manual / eager path: needed because F.scaled_dot_product_attention
            # is a fused kernel that never exposes the [B, H, T, T] softmax
            # probability matrix. This path is slower and uses more memory,
            # so it is only taken when weights are actually requested.
            scale = 1.0 / math.sqrt(self.head_dim)
            attn_scores = (q @ k.transpose(-2, -1)) * scale

            causal_mask = torch.full(
                (T, T), mask_value, device=x.device, dtype=torch.float32
            ).triu(diagonal=1)
            attn_scores = attn_scores + causal_mask.to(attn_scores.dtype)

            if attention_mask is not None:
                pad_bias = _expand_padding_mask(attention_mask, T, attn_scores.dtype, mask_value)
                if pad_bias is not None:
                    attn_scores = attn_scores + pad_bias

            attn_weights = F.softmax(attn_scores, dim=-1, dtype=torch.float32).to(q.dtype)
            # Guard against any residual NaN (e.g. a fully-padded row) so a
            # single degenerate row can't poison the whole loss.
            attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
            out = attn_weights @ v
        else:
            # Fast path: fused SDPA kernel, no materialized attention matrix.
            pad_bias = None
            if attention_mask is not None:
                pad_bias = _expand_padding_mask(attention_mask, T, q.dtype, mask_value)
            if pad_bias is not None:
                causal_mask = torch.full(
                    (T, T), mask_value, device=x.device, dtype=torch.float32
                ).triu(diagonal=1).to(q.dtype)
                combined_mask = causal_mask + pad_bias
                out = F.scaled_dot_product_attention(q, k, v, attn_mask=combined_mask)
            else:
                out = F.scaled_dot_product_attention(q, k, v, is_causal=True)

        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.o_proj(out), attn_weights


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

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


class Block(nn.Module):
    def __init__(self, config: IvmeXLConfig):
        super().__init__()
        self.attn_norm = RMSNorm(config.hidden_size, config.norm_eps)
        self.attn = Attention(config)
        self.mlp_norm = RMSNorm(config.hidden_size, config.norm_eps)
        self.mlp = SwiGLU(config.hidden_size, config.ffn_dim)

    def forward(
        self,
        x,
        cos,
        sin,
        attention_mask: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        attn_out, attn_weights = self.attn(
            self.attn_norm(x),
            cos,
            sin,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
        )
        x = x + attn_out
        x = x + self.mlp(self.mlp_norm(x))
        return x, attn_weights


class IvmeXLForCausalLM(PreTrainedModel):
    config_class = IvmeXLConfig
    base_model_prefix = "ivme_xl"
    _no_split_modules = ["Block"]
    supports_gradient_checkpointing = False

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

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList([Block(config) for _ in range(config.n_layers)])
        self.final_norm = RMSNorm(config.hidden_size, config.norm_eps)

        if config.tie_word_embeddings:
            self.lm_head = None
        else:
            self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        cos, sin = precompute_rope_freqs(
            config.head_dim, config.max_position_embeddings, config.rope_theta,
            device="cpu", dtype=torch.float32,
        )
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)

        self.post_init()

    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)

    def get_input_embeddings(self):
        return self.embed_tokens

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

    def get_output_embeddings(self):
        return self.embed_tokens if self.config.tie_word_embeddings else self.lm_head

    def set_output_embeddings(self, new_embeddings):
        if self.config.tie_word_embeddings:
            self.embed_tokens = new_embeddings
        else:
            self.lm_head = new_embeddings

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = True,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        output_attentions = (
            output_attentions
            if output_attentions is not None
            else getattr(self.config, "output_attentions", False)
        )

        B, T = input_ids.shape
        x = self.embed_tokens(input_ids)
        cos = self.rope_cos.to(device=x.device, dtype=x.dtype)
        sin = self.rope_sin.to(device=x.device, dtype=x.dtype)

        all_attentions = () if output_attentions else None

        for block in self.layers:
            x, attn_weights = block(
                x, cos, sin,
                attention_mask=attention_mask,
                output_attentions=output_attentions,
            )
            if output_attentions:
                all_attentions = all_attentions + (attn_weights,)

        x = self.final_norm(x)

        weight = self.embed_tokens.weight if self.config.tie_word_embeddings else self.lm_head.weight
        logits = F.linear(x, weight)

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous().to(logits.device)
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)).float(),
                shift_labels.view(-1).long(),
                ignore_index=-100,
            )

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            attentions=all_attentions,
        )

    @torch.no_grad()
    def generate(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int = 100,
        temperature: float = 1.0,
        top_k: Optional[int] = None,
        do_sample: bool = False,
        eos_token_id: Optional[int] = None,
        **kwargs,
    ) -> torch.Tensor:
        """
        Minimal standalone generation loop. transformers' full `.generate()`
        (with beam search, repetition penalty, etc.) will also work via the
        standard GenerationMixin machinery once this class is registered,
        since CausalLMOutputWithPast + a HF PreTrainedModel base class is
        exactly what GenerationMixin expects — but this simple override
        exists as a fast, dependency-light path that works even in contexts
        where the full generation utilities aren't wired up.
        """
        self.eval()
        max_len = self.config.max_position_embeddings

        for _ in range(max_new_tokens):
            ids_cond = input_ids[:, -max_len:]
            logits = self(ids_cond).logits
            next_logits = logits[:, -1, :] / max(temperature, 1e-5)

            if top_k is not None:
                v, _ = torch.topk(next_logits, min(top_k, next_logits.size(-1)))
                next_logits[next_logits < v[:, [-1]]] = -float("inf")

            if do_sample:
                probs = F.softmax(next_logits, dim=-1)
                next_id = torch.multinomial(probs, num_samples=1)
            else:
                next_id = next_logits.argmax(dim=-1, keepdim=True)

            input_ids = torch.cat([input_ids, next_id], dim=1)

            if eos_token_id is not None and (next_id == eos_token_id).all():
                break

        return input_ids


__all__ = ["IvmeXLConfig", "IvmeXLForCausalLM"]