"""
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.
"""

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


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) -> 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)

        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)


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):
        x = x + self.attn(self.attn_norm(x), cos, sin)
        x = x + self.mlp(self.mlp_norm(x))
        return x


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,
        return_dict: Optional[bool] = True,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        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)

        for block in self.layers:
            x = block(x, cos, sin)
        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()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)).float(),
                shift_labels.view(-1),
            )

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

    @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"]