"""
GPT model (Hugging Face `transformers`-compatible, self-contained).

A faithful, dependency-light port of the training-time GPT, shipped inside the
model repo and loaded via `trust_remote_code=True`:

    from transformers import AutoModelForCausalLM, AutoTokenizer
    model = AutoModelForCausalLM.from_pretrained("<repo>", trust_remote_code=True)
    tok   = AutoTokenizer.from_pretrained("<repo>", trust_remote_code=True)
    ids   = tok.build_chatml_ids([{"role": "user", "content": "Hello, who are you?"}])
    out   = model.chat_generate(ids)          # list[int] of generated token ids
    print(tok.decode_response(out))

Architecture: GPT-style decoder with RMSNorm (pre-norm), ReZero residual gates,
rotary position embeddings (RoPE) applied to Q/K, bias-free attention/MLP
Linears, and token/LM-head weight tying.

Differences from the training implementation (all inference-irrelevant):
  * NO `flash_attn` import. Attention uses `F.scaled_dot_product_attention`
    (the exact path training already takes for unpacked / inference), so this
    runs on a plain CPU box (e.g. a Hugging Face Space) with no CUDA build.
  * RMSNorm is a plain module (same float32-reduction forward numerics as the
    training custom-autograd version; the autograd version was only a training
    memory optimization).
  * Adds a KV-cached `stream_generate` / `chat_generate` so token-by-token
    generation on CPU is not O(n^2).

The submodule names (wte, h.<i>.{ln_1,attn.c_attn,attn.c_proj,ln_2,
mlp.c_fc,mlp.c_proj,attn_scale,mlp_scale}, ln_f, lm_head) match the exported
checkpoint's state_dict keys 1:1. `lm_head.weight` is tied to `wte.weight`.
"""

import math

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

from .configuration_gpt import GPTConfig


class RMSNorm(nn.Module):
    """Root-mean-square layer norm (no mean subtraction, no bias).

    Reduction is computed in float32 regardless of input dtype, matching the
    training-time RMSNorm forward exactly.
    """

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

    def forward(self, x):
        comp = x.dtype if x.dtype in (torch.float32, torch.float64) else torch.float32
        xf = x.to(comp)
        rstd = torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)
        return (xf * rstd * self.weight.to(comp)).to(x.dtype)


# ---- Rotary position embeddings (mirrors the training implementation) ----
def build_rope_cache(block_size, head_dim, theta):
    """cos/sin tables (block_size, head_dim); angles computed in float32."""
    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
    t = torch.arange(block_size, dtype=torch.float32)
    freqs = torch.outer(t, inv_freq)
    emb = torch.cat((freqs, freqs), dim=-1)
    return emb.cos(), emb.sin()


def rotate_half(x):
    half = x.shape[-1] // 2
    return torch.cat((-x[..., half:], x[..., :half]), dim=-1)


def apply_rope(q, k, cos, sin):
    """q, k: (B, T, nH, hd); cos/sin: (B or 1, T, 1, hd), gathered per token."""
    cos = cos.to(q.dtype)
    sin = sin.to(q.dtype)
    return q * cos + rotate_half(q) * sin, k * cos + rotate_half(k) * sin


class MultiHeadAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.head_dim = config.n_embd // config.n_head
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)

    def forward(self, x, rope, layer_past=None, use_cache=False):
        B, T, C = x.size()
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
        q = q.view(B, T, self.n_head, self.head_dim)
        k = k.view(B, T, self.n_head, self.head_dim)
        v = v.view(B, T, self.n_head, self.head_dim)

        # Rotate the CURRENT tokens' q/k by their absolute positions before
        # caching/concatenating — cached K entries are already rotated.
        q, k = apply_rope(q, k, rope[0], rope[1])

        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        if layer_past is not None:
            past_k, past_v = layer_past
            k = torch.cat((past_k, k), dim=2)
            v = torch.cat((past_v, v), dim=2)

        present = (k, v) if use_cache else None

        # Prefill (no past, T>1) needs the causal mask. A single cached decode
        # step (T==1) attends to all cached keys with no mask. The mixed case
        # (T>1 with a past) is not produced by our generation loop.
        is_causal = layer_past is None and T > 1
        y = F.scaled_dot_product_attention(q, k, v, attn_mask=None, is_causal=is_causal)
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.c_proj(y), present


class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)

    def forward(self, x):
        return self.c_proj(F.gelu(self.c_fc(x)))


class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = RMSNorm(config.n_embd, eps=config.rms_norm_eps)
        self.attn = MultiHeadAttention(config)
        self.ln_2 = RMSNorm(config.n_embd, eps=config.rms_norm_eps)
        self.mlp = MLP(config)
        # ReZero residual gates (one learned scalar per sublayer).
        self.attn_scale = nn.Parameter(torch.zeros(1))
        self.mlp_scale = nn.Parameter(torch.zeros(1))

    def forward(self, x, rope, layer_past=None, use_cache=False):
        a, present = self.attn(self.ln_1(x), rope, layer_past=layer_past, use_cache=use_cache)
        x = x + self.attn_scale * a
        x = x + self.mlp_scale * self.mlp(self.ln_2(x))
        return x, present


def _ban_repeat_ngrams(logits_row, gen_tokens, n):
    """Block any n-gram from repeating within the generated tokens (in place)."""
    if not n or n < 2 or len(gen_tokens) < n:
        return
    prefix = tuple(gen_tokens[-(n - 1):])
    for i in range(len(gen_tokens) - n + 1):
        if tuple(gen_tokens[i:i + n - 1]) == prefix:
            logits_row[gen_tokens[i + n - 1]] = float("-inf")


class GPTForCausalLM(PreTrainedModel):
    config_class = GPTConfig
    _no_split_modules = ["Block"]
    # Tells transformers that lm_head.weight is tied to the input embeddings, so
    # the loader does NOT treat it as a missing/freshly-initialized weight and
    # re-ties it after loading. Without this the head loads as random noise.
    # transformers >= 5.0 expects a dict {tied_key: source_key}.
    _tied_weights_keys = {"lm_head.weight": "wte.weight"}

    def __init__(self, config):
        super().__init__(config)
        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
        # RoPE cos/sin tables are NOT registered as buffers: transformers'
        # meta-device loading materializes non-persistent buffers as EMPTY
        # (garbage) memory, since they aren't in the safetensors and
        # __init__-computed values are discarded. A plain lazy attribute is
        # invisible to the loader — built from real math on first forward.
        self._rope_cache = None
        self.h = nn.ModuleList([Block(config) for _ in range(config.n_layer)])
        self.ln_f = RMSNorm(config.n_embd, eps=config.rms_norm_eps)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        if config.tie_weights:
            self.lm_head.weight = self.wte.weight
        self.post_init()

    def _rope_tables(self, device):
        """Lazy fp32 cos/sin tables on `device` (built once, rebuilt on device move)."""
        if self._rope_cache is None or self._rope_cache[0].device != device:
            cos, sin = build_rope_cache(
                self.config.block_size, self.config.n_embd // self.config.n_head,
                float(self.config.rope_theta))
            self._rope_cache = (cos.to(device), sin.to(device))
        return self._rope_cache

    # --- transformers plumbing (embeddings / weight tying) ---
    def get_input_embeddings(self):
        return self.wte

    def set_input_embeddings(self, new):
        self.wte = new

    def get_output_embeddings(self):
        return self.lm_head

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

    # --- standard forward: returns logits over the full sequence ---
    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        position_ids=None,
        past_key_values=None,
        use_cache=False,
        labels=None,
        return_dict=True,
        **kwargs,
    ):
        B, T = input_ids.size()
        device = input_ids.device

        past_len = past_key_values[0][0].size(2) if past_key_values is not None else 0
        if position_ids is None:
            position_ids = torch.arange(past_len, past_len + T, dtype=torch.long, device=device).unsqueeze(0)

        # Per-token cos/sin (B or 1, T, 1, head_dim): cached decode rotates the
        # single new token by its absolute position (past_len + i).
        cos_tab, sin_tab = self._rope_tables(device)
        rope = (cos_tab[position_ids].unsqueeze(2),
                sin_tab[position_ids].unsqueeze(2))

        x = self.wte(input_ids)

        presents = [] if use_cache else None
        for i, block in enumerate(self.h):
            layer_past = past_key_values[i] if past_key_values is not None else None
            x, present = block(x, rope, layer_past=layer_past, use_cache=use_cache)
            if use_cache:
                presents.append(present)

        x = self.ln_f(x)
        logits = self.lm_head(x)

        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)),
                shift_labels.view(-1),
                ignore_index=self.config.unk_pad_id,
            )

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

    # --- convenient KV-cached generation (used by the demo Space) ---
    @torch.no_grad()
    def stream_generate(
        self,
        input_ids,
        max_new_tokens=512,
        temperature=0.0,
        top_k=50,
        top_p=0.9,
        no_repeat_ngram_size=0,
    ):
        """Yield generated token ids one at a time (KV-cached).

        Mirrors the training-time chat decoding: bans <unk>/<pad> (id 0), applies a
        no-repeat-ngram block over the generated tokens only (off by default, 0: the
        trigram ban wrecked long answers, Wayne 2026-09-08), and stops on
        <|im_end|>. `input_ids` is a 1-D LongTensor of the ChatML prompt token ids.
        """
        device = next(self.parameters()).device
        idx = input_ids.to(device).view(1, -1)
        block_size = self.config.block_size
        if idx.size(1) >= block_size:
            idx = idx[:, -(block_size - 1):]

        eos_id = self.config.im_end_id
        ban_id = self.config.unk_pad_id

        past = None
        cur = idx
        generated = []
        for _ in range(max_new_tokens):
            if past is not None and past[0][0].size(2) >= block_size - 1:
                break  # would exceed the trained context window
            out = self(cur, past_key_values=past, use_cache=True)
            past = out.past_key_values
            logits = out.logits[0, -1, :].float()

            logits[ban_id] = float("-inf")
            _ban_repeat_ngrams(logits, generated, no_repeat_ngram_size)

            if temperature < 0.01:
                nxt = int(torch.argmax(logits).item())
            else:
                logits = logits / temperature
                if top_k:
                    v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                    logits[logits < v[-1]] = float("-inf")
                if top_p and top_p < 1.0:
                    sl, si = torch.sort(logits, descending=True)
                    cum = torch.cumsum(F.softmax(sl, dim=-1), dim=-1)
                    remove = cum > top_p
                    remove[1:] = remove[:-1].clone()
                    remove[0] = False
                    logits[si[remove]] = float("-inf")
                probs = F.softmax(logits, dim=-1)
                nxt = int(torch.multinomial(probs, 1).item())

            if nxt == eos_id:
                break
            generated.append(nxt)
            yield nxt
            cur = torch.tensor([[nxt]], dtype=torch.long, device=device)

    @torch.no_grad()
    def chat_generate(self, input_ids, **kwargs):
        """Non-streaming wrapper: returns the full list of generated token ids."""
        return list(self.stream_generate(input_ids, **kwargs))