"""TinyTransformer v2 as a Transformers custom model, servable by vLLM's Transformers modeling backend.

A copy of `MLExamples/TinyTransformer/version2_pytorch_fused/tiny_llama_v2.py`'s model from the https://github.com/amd/HPCTrainingExamples — except where v2's code
can't work when served. Every such change is marked `# SERVING CHANGE N:` below; there are three:
  1. Attention goes through `ALL_ATTENTION_FUNCTIONS` (keyed by `layer_idx`) instead of calling SDPA
     directly, so vLLM can plug in its own attention kernel and KV cache —
     https://docs.vllm.ai/en/latest/models/supported_models/#writing-custom-models. GQA's K/V head
     repeat moves into that interface too (transformers' `repeat_kv` gives the same head order as
     v2's `repeat_interleave`): vLLM sizes its KV cache by `n_kv_heads`.
  2. RoPE angles come from explicit `position_ids` instead of `start_pos + arange(seq_len)`, which
     assumes every call starts at position `start_pos` — wrong for vLLM's decoding steps and packed
     batches. vLLM may also hand in precomputed `inputs_embeds` instead of `input_ids`.
  3. transformers plumbing, so transformers/vLLM can load it: `PretrainedConfig`/`PreTrainedModel`
     base classes, and `TinyLlamaV2.forward` stops after the final norm and returns the hidden
     states. vLLM builds this class via `auto_map.AutoModel` and puts its own lm_head on top, tied to
     `get_input_embeddings()` (`tie_word_embeddings`) — which is exactly v2's `output_proj`, itself
     tied to `token_embedding`.

Also left out, without affecting outputs:
  - `record_function` profiler ranges and dropout (inactive at inference).
  - The unfused Q/K/V and gate/up projection branches, and the manual-attention/explicit-mask branch:
    the workshop checkpoint uses the fused layout and the SDPA path (`convert_and_push.py` refuses
    any other checkpoint), so `FusionConfig` is gone too.
  - `output_proj` and the loss (see 3), `get_fusion_statistics`, and `_init_weights`: random
    initialization, which loading the checkpoint overwrites.
"""

from typing import ClassVar

import torch
import torch.nn.functional as F
from torch import nn
from transformers import PretrainedConfig, PreTrainedModel
from transformers.integrations.sdpa_attention import sdpa_attention_forward
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS


class TinyLlamaConfig(PretrainedConfig):
    """v2's TinyLlamaConfig: same name, fields and defaults. SERVING CHANGE 3: a PretrainedConfig
    subclass instead of a dataclass, so transformers/vLLM can read it from config.json."""

    model_type = "tiny_llama_v2"
    # SERVING CHANGE 3: transformers/vLLM read the standard names; alias them to v2's fields.
    attribute_map: ClassVar[dict[str, str]] = {
        "hidden_size": "hidden_dim",
        "num_hidden_layers": "n_layers",
        "num_attention_heads": "n_heads",
        "num_key_value_heads": "n_kv_heads",
        "intermediate_size": "intermediate_dim",
        "rms_norm_eps": "norm_eps",
        "max_position_embeddings": "max_seq_len",
    }

    def __init__(
        self,
        vocab_size: int = 1000,
        hidden_dim: int = 256,
        n_layers: int = 4,
        n_heads: int = 8,
        n_kv_heads: int = 4,
        intermediate_dim: int = 512,
        max_seq_len: int = 128,
        rope_theta: float = 10000.0,
        norm_eps: float = 1e-6,
        dropout: float = 0.0,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.hidden_dim = hidden_dim
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.intermediate_dim = intermediate_dim
        self.max_seq_len = max_seq_len
        self.rope_theta = rope_theta
        self.norm_eps = norm_eps
        self.dropout = (
            dropout  # kept for fidelity with v2's config; unused at inference
        )
        # v2 always ties output_proj to token_embedding; vLLM reads this flag to tie its lm_head.
        kwargs["tie_word_embeddings"] = True
        super().__init__(**kwargs)


class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization - optimized for fusion."""

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        rms = torch.sqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
        x = x / rms
        return x * self.weight


class RotaryPositionEmbedding(nn.Module):
    """Rotary Position Embeddings, v2's interleaved-pair formula.

    SERVING CHANGE 2: angles are computed for the given `position_ids` rather than cached for
    `start_pos .. start_pos + seq_len`. v2 builds positions and frequencies in the activations'
    dtype (float32 for this checkpoint); so does this.
    """

    def __init__(self, dim: int, theta: float = 10000.0):
        super().__init__()
        self.dim = dim
        self.theta = theta

    def cos_sin(self, position_ids: torch.Tensor, dtype: torch.dtype):
        # position_ids: [batch, seq_len] -> cos/sin: [batch, seq_len, 1, dim // 2]
        t = position_ids.to(dtype)
        inv_freq = 1.0 / (
            self.theta
            ** (torch.arange(0, self.dim, 2, device=t.device, dtype=dtype) / self.dim)
        )
        freqs = t[..., None] * inv_freq
        return torch.cos(freqs).to(dtype).unsqueeze(2), torch.sin(freqs).to(
            dtype
        ).unsqueeze(2)

    @staticmethod
    def apply_rotary_embedding(
        q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
    ):
        batch_size, seq_len, n_heads_q, head_dim = q.shape
        _, _, n_heads_k, _ = k.shape

        # Reshape for rotary embedding
        q = q.reshape(batch_size, seq_len, n_heads_q, head_dim // 2, 2)
        k = k.reshape(batch_size, seq_len, n_heads_k, head_dim // 2, 2)

        q_rot = torch.stack(
            [q[..., 0] * cos - q[..., 1] * sin, q[..., 0] * sin + q[..., 1] * cos],
            dim=-1,
        )

        k_rot = torch.stack(
            [k[..., 0] * cos - k[..., 1] * sin, k[..., 0] * sin + k[..., 1] * cos],
            dim=-1,
        )

        # Reshape back
        q_rot = q_rot.reshape(batch_size, seq_len, n_heads_q, head_dim)
        k_rot = k_rot.reshape(batch_size, seq_len, n_heads_k, head_dim)

        return q_rot, k_rot


class FusedAttention(nn.Module):
    """Optimized attention with QKV fusion."""

    def __init__(self, config: TinyLlamaConfig, layer_idx: int):
        super().__init__()
        self.n_heads = config.n_heads
        self.n_kv_heads = config.n_kv_heads
        self.head_dim = config.hidden_dim // config.n_heads
        self.scale = self.head_dim**-0.5
        # SERVING CHANGE 1: what the attention interface reads off the module.
        self.config = config
        self.layer_idx = layer_idx
        self.num_key_value_groups = self.n_heads // self.n_kv_heads
        self.is_causal = True

        # Fused QKV projection - 3 operations combined into 1
        self.qkv_proj = nn.Linear(
            config.hidden_dim,
            (config.n_heads + 2 * config.n_kv_heads) * self.head_dim,
            bias=False,
        )
        self.o_proj = nn.Linear(
            config.n_heads * self.head_dim, config.hidden_dim, bias=False
        )

    def forward(
        self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
    ) -> torch.Tensor:
        batch_size, seq_len, _ = x.shape

        # Fused QKV projection, split into Q, K, V
        qkv = self.qkv_proj(x)
        q_size = self.n_heads * self.head_dim
        kv_size = self.n_kv_heads * self.head_dim

        q = qkv[:, :, :q_size].view(batch_size, seq_len, self.n_heads, self.head_dim)
        k = qkv[:, :, q_size : q_size + kv_size].view(
            batch_size, seq_len, self.n_kv_heads, self.head_dim
        )
        v = qkv[:, :, q_size + kv_size :].view(
            batch_size, seq_len, self.n_kv_heads, self.head_dim
        )

        # Apply rotary position embeddings
        q, k = RotaryPositionEmbedding.apply_rotary_embedding(q, k, cos, sin)

        # Transpose for attention computation
        q = q.transpose(1, 2)  # (batch, n_heads, seq_len, head_dim)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # SERVING CHANGE 1: v2 repeats the K/V heads here and calls
        #     F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True)
        # directly. transformers' "sdpa" interface is that same call, GQA repeat included, and lets
        # vLLM swap in its own kernel and KV cache. It returns [batch, seq_len, n_heads, head_dim],
        # already transposed back.
        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
            self.config._attn_implementation, sdpa_attention_forward
        )
        attn_output, _ = attention_interface(
            self, q, k, v, None, dropout=0.0, scaling=self.scale
        )

        attn_output = attn_output.contiguous().view(batch_size, seq_len, -1)
        return self.o_proj(attn_output)


class FusedSwiGLU(nn.Module):
    """Optimized SwiGLU with fused gate/up projections."""

    def __init__(self, config: TinyLlamaConfig):
        super().__init__()
        # Fused gate and up projection - 2 operations combined into 1
        self.gate_up_proj = nn.Linear(
            config.hidden_dim, 2 * config.intermediate_dim, bias=False
        )
        self.down_proj = nn.Linear(
            config.intermediate_dim, config.hidden_dim, bias=False
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate_up = self.gate_up_proj(x)
        gate, up = gate_up.chunk(2, dim=-1)
        intermediate = F.silu(gate) * up
        return self.down_proj(intermediate)


class FusedTransformerBlock(nn.Module):
    """Optimized transformer block with fusion capabilities."""

    def __init__(self, config: TinyLlamaConfig, layer_idx: int):
        super().__init__()
        self.attention = FusedAttention(config, layer_idx)
        self.feed_forward = FusedSwiGLU(config)
        self.norm1 = RMSNorm(config.hidden_dim, config.norm_eps)
        self.norm2 = RMSNorm(config.hidden_dim, config.norm_eps)

    def forward(
        self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
    ) -> torch.Tensor:
        # Pre-norm attention with residual connection
        x = x + self.attention(self.norm1(x), cos, sin)
        # Pre-norm feed-forward with residual connection
        x = x + self.feed_forward(self.norm2(x))
        return x


class TinyLlamaV2(PreTrainedModel):
    """Tiny LLaMA V2 with comprehensive fusion optimizations.

    SERVING CHANGE 3: a PreTrainedModel ending at the final norm (see module docstring).
    """

    config_class = TinyLlamaConfig
    _supports_sdpa = True
    # vLLM's Transformers backend requires this (or source it can inspect) to use its attention.
    _supports_attention_backend = True

    def __init__(self, config: TinyLlamaConfig):
        super().__init__(config)

        # Token embeddings
        self.token_embedding = nn.Embedding(config.vocab_size, config.hidden_dim)

        # Transformer blocks with fusion
        self.blocks = nn.ModuleList(
            [
                FusedTransformerBlock(config, layer_idx)
                for layer_idx in range(config.n_layers)
            ]
        )

        # Final norm
        self.norm = RMSNorm(config.hidden_dim, config.norm_eps)

        self.rope = RotaryPositionEmbedding(
            config.hidden_dim // config.n_heads, config.rope_theta
        )

        self.post_init()

    def get_input_embeddings(self):
        return self.token_embedding

    def forward(
        self,
        input_ids: torch.Tensor | None = None,
        position_ids: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
        **kwargs,
    ):
        # SERVING CHANGE 2: vLLM may pass embeddings, and always passes positions; default to v2's
        # 0..seq_len-1. vLLM also passes use_cache/return_dict etc., unused here (**kwargs).
        x = self.token_embedding(input_ids) if inputs_embeds is None else inputs_embeds
        if position_ids is None:
            position_ids = torch.arange(x.shape[1], device=x.device).unsqueeze(0)
        cos, sin = self.rope.cos_sin(position_ids, x.dtype)

        for block in self.blocks:
            x = block(x, cos, sin)

        # SERVING CHANGE 3: stop at the final norm; vLLM applies its tied lm_head. A tuple — vLLM
        # reads the hidden states as outputs[0].
        return (self.norm(x),)