"""TinyBrainBot recurrent-depth model for transformers (trust_remote_code=True).

Below the marker is the exact architecture source the model was trained with. The wrapper builds it from the
config, so weights load unchanged. Generation recomputes the prefix each step (no KV cache): fine at this size.
"""
from __future__ import annotations
# ======================= native architecture (verbatim) =======================
"""TinyBrainBot reference architecture.

Matches the blueprint spec:
    - decoder-only transformer
    - 8 layers, hidden 1536, 12 heads (head_dim = 128), FFN 6144
    - RoPE positional encoding
    - RMSNorm (pre-norm)
    - SwiGLU feed-forward
    - 36k Unigram vocab
    - ctx 2048

The defaults here produce a ~300M parameter model. Tests use smaller overrides.
"""

from dataclasses import dataclass, field, asdict
from typing import Optional

import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint as _ckpt


@dataclass
class TinyBrainBotConfig:
    # Blueprint defaults
    vocab_size: int = 36_000
    hidden_size: int = 1536
    n_layers: int = 8
    n_heads: int = 12
    head_dim: int = 128
    n_kv_heads: Optional[int] = None   # None => MHA (== n_heads); a smaller divisor => GQA
    ffn_size: int = 6144
    context_length: int = 2048
    rope_theta: float = 10_000.0
    rms_norm_eps: float = 1e-5
    tie_embeddings: bool = True
    init_std: float = 0.02
    dropout: float = 0.0
    gradient_checkpointing: bool = False
    # recurrent depth (opt-in): store `loop_unique` blocks and run them `n_loops` times. 0 = ordinary stack.
    loop_unique: int = 0
    n_loops: int = 1
    loop_train_min: int = 0          # training samples the loop count from [min, max] per step (0 = n_loops)
    loop_train_max: int = 0
    # digit-place embedding (opt-in): each digit token gets its index in its digit run (1 = first digit; with
    # units-first numbers that is the units digit). Digit token ids are digit_lo..digit_hi.
    digit_place: bool = False
    digit_lo: int = -1
    digit_hi: int = -1
    place_vocab: int = 16
    # factorized embedding (opt-in): a vocab x embed_rank table, projected up to hidden_size on the way in and down
    # to embed_rank before the (tied) output layer. 0 = ordinary vocab x hidden_size embedding.
    embed_rank: int = 0

    @property
    def attention_dim(self) -> int:
        return self.n_heads * self.head_dim

    @property
    def num_kv_heads(self) -> int:
        """KV head count — equals n_heads for MHA, fewer for GQA."""
        return self.n_kv_heads if self.n_kv_heads is not None else self.n_heads

    @property
    def kv_dim(self) -> int:
        return self.num_kv_heads * self.head_dim

    def summary(self) -> dict:
        return {
            **asdict(self),
            "attention_dim": self.attention_dim,
        }


# ----------------------------------------------------------------------------
# Building blocks
# ----------------------------------------------------------------------------

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:
        # computed in float32 for stability, cast back to input dtype
        orig_dtype = x.dtype
        x = x.float()
        variance = x.pow(2).mean(dim=-1, keepdim=True)
        x = x * torch.rsqrt(variance + self.eps)
        return (self.weight * x).to(orig_dtype)


def _precompute_rope_cache(dim: int, max_len: int, theta: float, device, dtype):
    # standard RoPE frequencies
    inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, device=device).float() / dim))
    t = torch.arange(max_len, device=device).float()
    freqs = torch.outer(t, inv_freq)  # [max_len, dim/2]
    cos = freqs.cos().to(dtype)
    sin = freqs.sin().to(dtype)
    return cos, sin


def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    # x: [B, H, T, D]; cos/sin: [T, D/2], or [B, T, D/2] for per-row positions (padded batches)
    x1, x2 = x[..., ::2], x[..., 1::2]
    if cos.dim() == 3:
        cos_b, sin_b = cos.unsqueeze(1), sin.unsqueeze(1)
    else:
        # broadcast cos/sin over [B, H, T, D/2]
        cos_b = cos.unsqueeze(0).unsqueeze(0)
        sin_b = sin.unsqueeze(0).unsqueeze(0)
    y1 = x1 * cos_b - x2 * sin_b
    y2 = x1 * sin_b + x2 * cos_b
    out = torch.stack([y1, y2], dim=-1).flatten(-2)
    return out


def _repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    """Expand KV heads to match query heads for GQA. x: [B, n_kv, T, head_dim]."""
    if n_rep == 1:
        return x
    B, n_kv, T, hd = x.shape
    return x[:, :, None, :, :].expand(B, n_kv, n_rep, T, hd).reshape(B, n_kv * n_rep, T, hd)


class CausalSelfAttention(nn.Module):
    def __init__(self, cfg: TinyBrainBotConfig):
        super().__init__()
        self.cfg = cfg
        nkv = cfg.num_kv_heads
        assert cfg.n_heads % nkv == 0, "n_heads must be divisible by n_kv_heads"
        # Combined Q/K/V. With MHA (nkv == n_heads) this is 3*attention_dim — the exact
        # shape older checkpoints stored, so they load unchanged. With GQA it shrinks K/V.
        self.qkv = nn.Linear(cfg.hidden_size, (cfg.n_heads + 2 * nkv) * cfg.head_dim, bias=False)
        self.out_proj = nn.Linear(cfg.attention_dim, cfg.hidden_size, bias=False)
        self.scale = 1.0 / math.sqrt(cfg.head_dim)
        self._q_size = cfg.n_heads * cfg.head_dim
        self._kv_size = nkv * cfg.head_dim
        self._n_rep = cfg.n_heads // nkv

    def forward(
        self,
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        past_kv: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        attn_mask: Optional[torch.Tensor] = None,
    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        B, T, _ = x.shape
        nkv = self.cfg.num_kv_heads
        q, k, v = self.qkv(x).split([self._q_size, self._kv_size, self._kv_size], dim=-1)
        q = q.view(B, T, self.cfg.n_heads, self.cfg.head_dim).transpose(1, 2)
        k = k.view(B, T, nkv, self.cfg.head_dim).transpose(1, 2)
        v = v.view(B, T, nkv, self.cfg.head_dim).transpose(1, 2)

        if cos.dim() == 3:                       # per-row positions, already [B, T, D/2]
            q = _apply_rope(q, cos, sin)
            k = _apply_rope(k, cos, sin)
        else:
            offset = past_kv[0].shape[-2] if past_kv is not None else 0
            q = _apply_rope(q, cos[offset:offset + T], sin[offset:offset + T])
            k = _apply_rope(k, cos[offset:offset + T], sin[offset:offset + T])

        if past_kv is not None:
            k = torch.cat([past_kv[0], k], dim=-2)
            v = torch.cat([past_kv[1], v], dim=-2)
        new_kv = (k, v)  # cache holds the SMALL (n_kv) heads — the whole point of GQA

        # Expand KV to n_heads for attention (no-op under MHA).
        k = _repeat_kv(k, self._n_rep)
        v = _repeat_kv(v, self._n_rep)

        if attn_mask is not None:                # padded batch: causal AND not-a-pad-key, [B, 1, T, T] bool
            attn = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
        else:
            is_causal = T > 1
            attn = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
        attn = attn.transpose(1, 2).contiguous().view(B, T, self.cfg.attention_dim)
        return self.out_proj(attn), new_kv


class SwiGLU(nn.Module):
    """Gated SiLU feed-forward (as in LLaMA / Mistral)."""

    def __init__(self, cfg: TinyBrainBotConfig):
        super().__init__()
        self.gate = nn.Linear(cfg.hidden_size, cfg.ffn_size, bias=False)
        self.up = nn.Linear(cfg.hidden_size, cfg.ffn_size, bias=False)
        self.down = nn.Linear(cfg.ffn_size, cfg.hidden_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down(F.silu(self.gate(x)) * self.up(x))


class DecoderBlock(nn.Module):
    def __init__(self, cfg: TinyBrainBotConfig):
        super().__init__()
        self.attn_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.attn = CausalSelfAttention(cfg)
        self.ffn_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.ffn = SwiGLU(cfg)

    def forward(
        self,
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        past_kv: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        attn_mask: Optional[torch.Tensor] = None,
    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        attn_out, new_kv = self.attn(self.attn_norm(x), cos, sin, past_kv, attn_mask)
        x = x + attn_out
        x = x + self.ffn(self.ffn_norm(x))
        return x, new_kv

    def forward_no_cache(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        out, _ = self.forward(x, cos, sin, None)
        return out


class TinyBrainBotModel(nn.Module):
    def __init__(self, cfg: TinyBrainBotConfig):
        super().__init__()
        self.cfg = cfg
        r = cfg.embed_rank if cfg.embed_rank > 0 else cfg.hidden_size
        self.embed = nn.Embedding(cfg.vocab_size, r)
        self.embed_up = nn.Linear(r, cfg.hidden_size, bias=False) if cfg.embed_rank > 0 else None
        self.head_down = nn.Linear(cfg.hidden_size, r, bias=False) if cfg.embed_rank > 0 else None
        n_blocks = cfg.loop_unique if cfg.loop_unique > 0 else cfg.n_layers
        self.blocks = nn.ModuleList([DecoderBlock(cfg) for _ in range(n_blocks)])
        # recurrent depth: a learned vector per loop pass, and a gate re-injecting the input embedding each pass
        # (named *_norm so the optimiser puts it in the no-weight-decay group)
        self.loop_emb = nn.Embedding(32, cfg.hidden_size) if cfg.loop_unique > 0 else None
        self.reinject_norm = nn.Parameter(torch.tensor(0.5)) if cfg.loop_unique > 0 else None
        self.place = nn.Embedding(cfg.place_vocab, cfg.hidden_size) if cfg.digit_place else None
        self.train_loops: Optional[int] = None     # set per step by the trainer (same on every rank)
        self.eval_loops: Optional[int] = None      # override for evaluation ("think longer")
        self.final_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.lm_head = nn.Linear(r, cfg.vocab_size, bias=False)
        if cfg.tie_embeddings:
            self.lm_head.weight = self.embed.weight

        self._rope_cache: Optional[tuple[torch.Tensor, torch.Tensor]] = None
        self._rope_device: Optional[torch.device] = None
        self._rope_dtype: Optional[torch.dtype] = None

        self.apply(self._init_weights)
        if self.embed_up is not None:
            # keep the residual stream at the usual scale: up-projection ~ 1/sqrt(rank)
            nn.init.normal_(self.embed_up.weight, std=r ** -0.5)
            nn.init.normal_(self.head_down.weight, std=cfg.hidden_size ** -0.5)

    def _init_weights(self, m: nn.Module) -> None:
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, std=self.cfg.init_std)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, std=self.cfg.init_std)

    def _place_ids(self, ids: torch.Tensor) -> torch.Tensor:
        d = (ids >= self.cfg.digit_lo) & (ids <= self.cfg.digit_hi)
        idx = torch.arange(ids.shape[1], device=ids.device).expand_as(ids)
        last_nd = torch.cummax(torch.where(d, torch.full_like(idx, -1), idx), dim=1).values
        return torch.where(d, (idx - last_nd).clamp(max=self.cfg.place_vocab - 1), torch.zeros_like(idx))

    def loops_now(self) -> int:
        if self.training and self.train_loops is not None:
            return self.train_loops
        if not self.training and self.eval_loops is not None:
            return self.eval_loops
        return self.cfg.n_loops

    def _ensure_rope(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]:
        need_rebuild = (
            self._rope_cache is None
            or self._rope_device != device
            or self._rope_dtype != dtype
            or self._rope_cache[0].shape[0] < seq_len
        )
        if need_rebuild:
            target_len = max(seq_len, self.cfg.context_length)
            cos, sin = _precompute_rope_cache(self.cfg.head_dim, target_len, self.cfg.rope_theta, device, dtype)
            self._rope_cache = (cos, sin)
            self._rope_device = device
            self._rope_dtype = dtype
        return self._rope_cache

    def forward(
        self,
        input_ids: torch.Tensor,
        labels: Optional[torch.Tensor] = None,
        past_kv: Optional[list[tuple[torch.Tensor, torch.Tensor]]] = None,
        use_cache: bool = False,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> dict:
        B, T = input_ids.shape
        past_len = past_kv[0][0].shape[-2] if past_kv is not None else 0
        if T + past_len > self.cfg.context_length:
            raise ValueError(f"sequence length {T + past_len} exceeds context_length {self.cfg.context_length}")
        x = self.embed(input_ids)
        if self.embed_up is not None:
            x = self.embed_up(x)
        if self.place is not None:
            x = x + self.place(self._place_ids(input_ids))
        cos, sin = self._ensure_rope(T + past_len, x.device, x.dtype)
        amask = None
        if attention_mask is not None and not bool(attention_mask.all()):
            # padded batch (left or right): RoPE positions count from each row's first real token, pad keys are
            # masked out, and pad query rows attend to themselves only (no NaN). Unpadded input takes the old path.
            if past_kv is not None or use_cache:
                raise ValueError("attention_mask with padding is not supported together with a KV cache")
            am = attention_mask.to(device=x.device, dtype=torch.bool)
            pos = (am.long().cumsum(-1) - 1).clamp(min=0)
            cos, sin = cos[pos], sin[pos]
            eye = torch.eye(T, dtype=torch.bool, device=x.device)
            amask = (eye.tril() | torch.ones(T, T, dtype=torch.bool, device=x.device).tril()) & am[:, None, None, :]
            amask = amask | eye[None, None]
        use_ckpt = self.cfg.gradient_checkpointing and self.training
        new_past: list[tuple[torch.Tensor, torch.Tensor]] = []
        if self.cfg.loop_unique > 0:
            if past_kv is not None or use_cache:
                raise ValueError("recurrent-depth models do not use a KV cache (recompute the prefix)")
            x0 = x
            for k in range(self.loops_now()):
                x = x + self.loop_emb.weight[k].to(x.dtype)
                if k > 0:
                    x = x + self.reinject_norm.to(x.dtype) * x0
                for block in self.blocks:
                    x = _ckpt(block.forward_no_cache, x, cos, sin, use_reentrant=False) if use_ckpt else block(x, cos, sin, None, amask)[0]
        for i, block in enumerate(self.blocks if self.cfg.loop_unique == 0 else []):
            layer_past = past_kv[i] if past_kv is not None else None
            if use_ckpt:
                x = _ckpt(block.forward_no_cache, x, cos, sin, use_reentrant=False)
            else:
                x, kv = block(x, cos, sin, layer_past, amask)
                if use_cache:
                    new_past.append(kv)
        x = self.final_norm(x)
        logits = self.lm_head(self.head_down(x) if self.head_down is not None else x)

        out = {"logits": logits}
        if use_cache and not use_ckpt:
            out["past_kv"] = new_past
        if labels is not None:
            # standard next-token loss, labels shifted upstream
            # unsqueeze(0) makes loss 1-D so torch.nn.DataParallel can
            # gather per-device losses along dim 0 (noop on single GPU).
            if (labels != -100).any():
                loss = F.cross_entropy(
                    logits.view(-1, self.cfg.vocab_size),
                    labels.view(-1),
                    ignore_index=-100,
                )
            else:
                # SFT edge case: a packed window can land entirely in masked
                # (user/system) tokens, leaving no supervised targets. cross_entropy
                # would then divide by zero and return NaN, poisoning training.
                # Emit a graph-connected zero loss for these batches instead.
                loss = logits.sum() * 0.0
            out["loss"] = loss.unsqueeze(0)
        return out

    def num_parameters(self) -> int:
        return sum(p.numel() for p in self.parameters())

# ======================= transformers wrapper =======================
from transformers import PreTrainedModel, GenerationMixin
from transformers.modeling_outputs import CausalLMOutput
from .configuration_tbbloop import TBBLoopConfig


class TBBLoopForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = TBBLoopConfig
    # transformers 5 wants {tied param: source param}; 4.x wants a list of tied params
    _tied_weights_keys = ({"model.lm_head.weight": "model.embed.weight"}
                          if int(__import__("transformers").__version__.split(".")[0]) >= 5 else ["model.lm_head.weight"])
    _supports_cache_class = False

    def __init__(self, config):
        super().__init__(config)
        self.model = TinyBrainBotModel(TinyBrainBotConfig(**config.native))
        if getattr(config, "eval_loops", 0):
            self.model.eval_loops = int(config.eval_loops)
        self.post_init()                         # required by transformers 5 (tied-weight bookkeeping); harmless on 4.x

    def _init_weights(self, module):
        pass                                     # weights always come from the checkpoint

    def get_input_embeddings(self):
        return self.model.embed

    def set_input_embeddings(self, value):
        self.model.embed = value

    def get_output_embeddings(self):
        return self.model.lm_head

    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
        logits = self.model(input_ids, attention_mask=attention_mask)["logits"]
        loss = None
        if labels is not None:
            loss = F.cross_entropy(logits[:, :-1].reshape(-1, logits.shape[-1]).float(), labels[:, 1:].reshape(-1), ignore_index=-100)
        return CausalLMOutput(loss=loss, logits=logits)

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {"input_ids": input_ids[:, -self.model.cfg.context_length:]}