# Copyright 2026. Released under the MIT license.
"""Self-contained looped decoder-only transformer for character-level arithmetic.

This module intentionally has no dependency on the training repository so it can
run from the Hub with ``trust_remote_code=True``. The building blocks mirror the
original ``model.py`` (RMSNorm, RoPE, causal attention, GELU MLP) and the
parameter names match the released checkpoint so weights load directly.
"""

from __future__ import annotations

import math
import re

import torch
import torch.nn as nn

from transformers import PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast

from .config import ThreeDigitBasicCalcConfig


class WordEmbeddings(nn.Module):
    def __init__(self, d_model: int, vocab_size: int):
        super().__init__()
        self.d_model = d_model
        self.embedding = nn.Embedding(vocab_size, d_model)

    def forward(self, x):
        return self.embedding(x) * math.sqrt(self.d_model)


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

    def forward(self, x):
        mean = x.pow(2).mean(dim=-1, keepdim=True)
        return self.weight * x * torch.rsqrt(mean + self.eps)


class RotaryEmbedding(nn.Module):
    def __init__(self, max_seq_len: int, h_dim: int, base: float = 10000.0):
        super().__init__()
        assert h_dim % 2 == 0
        self.max_seq_len = max_seq_len
        # Persistent so from_pretrained (low-mem loading) always materializes it;
        # cos/sin are derived per forward to avoid stale non-persistent buffers.
        inv_freq = 1.0 / (base ** (torch.arange(0, h_dim, 2).float() / h_dim))
        self.register_buffer("inv_freq", inv_freq, persistent=True)
        # Derived runtime values are deliberately not registered buffers: the
        # checkpoint contains only inv_freq, while repeated loop passes reuse
        # the trigonometric table.
        self._cos_cached = None
        self._sin_cached = None
        self._cache_key = None

    def _cos_sin(self, length: int, device):
        key = (length, device.type, device.index, self.inv_freq.dtype)
        if self._cache_key == key:
            return self._cos_cached, self._sin_cached
        t = torch.arange(length, device=device, dtype=self.inv_freq.dtype)
        freqs = torch.outer(t, self.inv_freq)
        self._cos_cached = freqs.cos()
        self._sin_cached = freqs.sin()
        self._cache_key = key
        return self._cos_cached, self._sin_cached

    def forward(self, x, position_ids=None):
        seq = x.shape[-2]
        if position_ids is None:
            position_ids = torch.arange(seq, device=x.device).unsqueeze(0)
        position_ids = position_ids.to(device=x.device, dtype=torch.long)
        required = int(position_ids.max().item()) + 1
        if required > self.max_seq_len:
            raise ValueError(
                f"position {required - 1} exceeds the configured context "
                f"length {self.max_seq_len}"
            )
        cos_table, sin_table = self._cos_sin(required, x.device)
        cos = cos_table[position_ids].unsqueeze(1)
        sin = sin_table[position_ids].unsqueeze(1)
        x1 = x[..., 0::2]
        x2 = x[..., 1::2]
        out1 = x1 * cos - x2 * sin
        out2 = x1 * sin + x2 * cos
        return torch.stack([out1, out2], dim=-1).flatten(-2).type_as(x)


class MultiHeadAttentionBlock(nn.Module):
    def __init__(self, max_seq_len, d_model, h, dropout=0.0, base=10000.0):
        super().__init__()
        assert d_model % h == 0
        self.heads = h
        self.d_head = d_model // h
        self.w_q = nn.Linear(d_model, d_model, bias=False)
        self.w_k = nn.Linear(d_model, d_model, bias=False)
        self.w_v = nn.Linear(d_model, d_model, bias=False)
        self.w_o = nn.Linear(d_model, d_model, bias=False)
        self.rope = RotaryEmbedding(max_seq_len, self.d_head, base)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None, position_ids=None):
        b, s, d = x.shape
        q = self.w_q(x).view(b, s, self.heads, self.d_head).transpose(1, 2)
        k = self.w_k(x).view(b, s, self.heads, self.d_head).transpose(1, 2)
        v = self.w_v(x).view(b, s, self.heads, self.d_head).transpose(1, 2)
        q = self.rope(q, position_ids)
        k = self.rope(k, position_ids)
        attn_score = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_head)
        if mask is not None:
            attn_score = attn_score.masked_fill(~mask, float("-inf"))
            # Left-padded query rows before the first real token have no legal
            # key. Their output is discarded; zeroing only those score rows
            # avoids NaNs while leaving every real-token score bit-identical to
            # the unpadded/native implementation.
            no_legal_key = ~mask.any(dim=-1, keepdim=True)
            attn_score = attn_score.masked_fill(no_legal_key, 0.0)
        attn = torch.softmax(attn_score, dim=-1)
        attn = self.dropout(attn)
        out = (attn @ v).transpose(1, 2).contiguous().reshape(b, s, d)
        return self.w_o(out)


class MLP(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.0):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        self.fc1 = nn.Linear(d_model, d_ff)
        self.fc2 = nn.Linear(d_ff, d_model)
        self.act = nn.GELU()

    def forward(self, x):
        return self.fc2(self.dropout(self.act(self.fc1(x))))


class DecoderBlock(nn.Module):
    def __init__(self, max_seq_len, d_model, attention_heads, d_ff,
                 eps=1e-6, dropout=0.0, base=10000.0):
        super().__init__()
        self.layer_norm_1 = RMSLayerNormalization(d_model, eps)
        self.layer_norm_2 = RMSLayerNormalization(d_model, eps)
        self.attn = MultiHeadAttentionBlock(
            max_seq_len, d_model, attention_heads, dropout, base)
        self.ffn = MLP(d_model, d_ff, dropout)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask, position_ids=None):
        x = x + self.dropout(
            self.attn(self.layer_norm_1(x), mask, position_ids)
        )
        x = x + self.dropout(self.ffn(self.layer_norm_2(x)))
        return x


class ThreeDigitBasicCalcForCausalLM(PreTrainedModel, GenerationMixin):
    """Looped decoder-only causal LM. Parameter names match the checkpoint."""

    config_class = ThreeDigitBasicCalcConfig
    _supports_cache_class = False

    def __init__(self, config: ThreeDigitBasicCalcConfig):
        super().__init__(config)
        d = config.hidden_size
        self.embedding = WordEmbeddings(d, config.vocab_size)
        self.decoder_layers = nn.ModuleList([
            DecoderBlock(config.max_position_embeddings, d,
                         config.num_attention_heads, 4 * d,
                         config.rms_eps, config.dropout, config.rope_base)
            for _ in range(config.num_hidden_layers)
        ])
        self.step_emb = nn.Embedding(config.n_loops, d)
        self.layer_norm = RMSLayerNormalization(d, config.rms_eps)
        self.projection = nn.Linear(d, config.vocab_size, bias=False)
        self.config.use_cache = False
        self.post_init()

    def get_input_embeddings(self):
        return self.embedding.embedding

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

    def get_output_embeddings(self):
        return self.projection

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=0.02)

    def forward(
        self,
        input_ids,
        attention_mask=None,
        position_ids=None,
        labels=None,
        **kwargs,
    ):
        b, s = input_ids.shape
        if s > self.config.max_position_embeddings:
            raise ValueError(
                f"sequence length {s} exceeds the configured context length "
                f"{self.config.max_position_embeddings}"
            )
        if attention_mask is None:
            attention_mask = torch.ones_like(input_ids, dtype=torch.long)
        if attention_mask.shape != input_ids.shape:
            raise ValueError("attention_mask must have the same shape as input_ids")
        attention_mask = attention_mask.to(device=input_ids.device)
        if position_ids is None:
            position_ids = attention_mask.long().cumsum(dim=-1) - 1
            position_ids = position_ids.clamp_min(0)
        if position_ids.shape[-1] != s:
            raise ValueError("position_ids must have the same sequence length as input_ids")

        causal = torch.tril(
            torch.ones(s, s, dtype=torch.bool, device=input_ids.device)
        ).view(1, 1, s, s)
        key_mask = attention_mask.to(torch.bool).view(b, 1, 1, s)
        combined_mask = causal & key_mask
        x = self.embedding(input_ids)
        for t in range(self.config.n_loops):
            x = x + self.step_emb.weight[t]
            for decoder in self.decoder_layers:
                x = decoder(x, combined_mask, position_ids)
        x = self.layer_norm(x)
        logits = self.projection(x)

        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            shift_attention = attention_mask[:, 1:].to(torch.bool)
            shift_labels = shift_labels.masked_fill(
                ~shift_attention | (shift_labels == self.config.pad_token_id),
                -100,
            )
            loss = nn.functional.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
            )
        return CausalLMOutputWithPast(loss=loss, logits=logits)

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        # No KV cache: the model recomputes the full sequence each step.
        attention_mask = kwargs.get("attention_mask")
        if attention_mask is None:
            attention_mask = torch.ones_like(input_ids, dtype=torch.long)
        position_ids = attention_mask.long().cumsum(dim=-1) - 1
        position_ids = position_ids.clamp_min(0)
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "position_ids": position_ids,
        }

    @torch.no_grad()
    def solve(self, tokenizer, expression: str, max_new_tokens: int = 256,
              return_trace: bool = False):
        """Compute a 3-digit arithmetic expression, e.g. ``model.solve(tok, "842/37")``.

        Returns the human-readable answer string (or a ``(answer, trace)`` tuple
        when ``return_trace=True``). The model emits an explicit scratchpad; this
        decodes its final ``<ans>`` field back to an ordinary number.
        """
        a, op, b, prompt = _parse_expression(expression)
        encoded = tokenizer(prompt, return_tensors="pt")
        ids = encoded.input_ids.to(self.device)
        remaining = self.config.max_position_embeddings - ids.shape[1]
        if remaining <= 0:
            raise ValueError("the expression leaves no context for generation")
        max_new_tokens = min(max_new_tokens, remaining)
        out = self.generate(
            ids,
            attention_mask=encoded.attention_mask.to(self.device),
            max_new_tokens=max_new_tokens,
            do_sample=False,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.pad_token_id,
        )
        if out[0, -1].item() != tokenizer.eos_token_id:
            raise ValueError(
                "generation reached the context limit before emitting <eos>"
            )
        trace = tokenizer.decode(out[0], skip_special_tokens=False)
        answer = _decode_answer(
            prompt,
            tokenizer.decode(out[0], skip_special_tokens=True),
            parsed_expression=(a, op, b),
        )
        return (answer, trace) if return_trace else answer


_EXPRESSION_RE = re.compile(
    r"^\s*(-?\d{1,3})\s*([+\-*/])\s*(-?\d{1,3})\s*=?\s*$"
)


def _parse_expression(expression: str):
    if not isinstance(expression, str):
        raise TypeError("expression must be a string")
    match = _EXPRESSION_RE.fullmatch(expression)
    if match is None:
        raise ValueError(
            "expected '<integer><operator><integer>' with operands from "
            "-999 through 999"
        )
    a, op, b = int(match.group(1)), match.group(2), int(match.group(3))
    prompt = f"{a}{op}{b}="
    return a, op, b, prompt


def _magnitude_plan(a: int, op: str, b: int):
    if op in ("*", "/"):
        effective = "mul" if op == "*" else "div"
        if op == "/" and b == 0:
            return effective, "positive", abs(a), 0
        sign = "negative" if (a < 0) != (b < 0) else "positive"
        return effective, sign, abs(a), abs(b)

    left_term = a
    right_term = b if op == "+" else -b
    if left_term == 0 and right_term == 0:
        return "add", "positive", 0, 0
    if left_term == 0:
        sign = "negative" if right_term < 0 else "positive"
        return "add", sign, 0, abs(right_term)
    if right_term == 0:
        sign = "negative" if left_term < 0 else "positive"
        return "add", sign, abs(left_term), 0
    if (left_term < 0) == (right_term < 0):
        sign = "negative" if left_term < 0 else "positive"
        return "add", sign, abs(left_term), abs(right_term)
    if abs(left_term) == abs(right_term):
        return "sub", "positive", abs(left_term), abs(right_term)
    if abs(left_term) > abs(right_term):
        sign = "negative" if left_term < 0 else "positive"
        return "sub", sign, abs(left_term), abs(right_term)
    sign = "negative" if right_term < 0 else "positive"
    return "sub", sign, abs(right_term), abs(left_term)


def _decode_answer(
    prompt: str,
    decoded: str,
    *,
    parsed_expression=None,
) -> str:
    """Strictly validate a v3 scratchpad and decode its fixed answer field."""
    if parsed_expression is None:
        a, op, b, canonical_prompt = _parse_expression(prompt)
    else:
        a, op, b = parsed_expression
        canonical_prompt = f"{a}{op}{b}="
    if prompt != canonical_prompt:
        raise ValueError("prompt is not in canonical form")
    if not decoded.startswith(canonical_prompt):
        raise ValueError("generated scratchpad does not begin with the prompt")

    effective, sign, left, right = _magnitude_plan(a, op, b)
    operation_token = {
        "add": "<add>",
        "sub": "<sub>",
        "mul": "<mul>",
        "div": "<div>",
    }[effective]
    sign_token = "<neg>" if sign == "negative" else "<pos>"
    if effective == "div":
        state = f"{left:03d}{right:03d}"
    else:
        state = f"{left:03d}"[::-1] + f"{right:03d}"[::-1]
    prefix = (
        re.escape(canonical_prompt)
        + re.escape(operation_token + sign_token + "<state>" + state)
    )

    if effective in ("add", "sub"):
        body = r"(?:<step>\d{5}){3}<ans>(?P<answer>\d{4})"
    elif effective == "mul":
        groups = []
        for term_count in (1, 2, 3, 2, 1):
            groups.append(rf"(?:<step>\d{{7}}){{{term_count}}}<col>\d{{3}}")
        body = "".join(groups) + r"<ans>(?P<answer>\d{6})"
    elif b == 0:
        body = r"<ans>(?P<answer><nan>)"
    else:
        record = r"<step>\d{8}<qmul>\d{5}<rem>\d{3}"
        body = rf"(?:{record}){{7}}<ans>(?P<answer>\d{{3}}\.\d{{3}})"

    match = re.fullmatch(prefix + body, decoded)
    if match is None:
        raise ValueError("generated scratchpad does not match the v3 grammar")
    internal = match.group("answer")
    if internal == "<nan>":
        return "NAN"
    if op == "/":
        integer, fraction = internal.split(".", 1)
        fraction = fraction.rstrip("0")
        magnitude = str(int(integer)) + (f".{fraction}" if fraction else "")
    else:
        magnitude = str(int(internal[::-1]))
    if sign == "negative" and magnitude != "0":
        return "-" + magnitude
    return magnitude