"""Decoder-only transformer for the IGLA-Coder pilot.

One architecture, two arms that differ only in the linear layers inside the
blocks:
  * fp      -- ordinary bf16/fp32 linears;
  * ternary -- BitNet b1.58 linears: parameter-free RMSNorm on the input,
               8-bit absmax activations per token and ternary {-1, 0, +1}
               absmean weights, both through straight-through estimators.
Embeddings, the tied output head and the norms stay full precision in both arms.

Layout follows the IGLA-Coder spec (t27 specs/igla/coder/arch.t27): RMSNorm,
RoPE, grouped-query attention, SwiGLU, tied embeddings (t27 #1034 P1.1).
"""
from dataclasses import dataclass, asdict
import math

import torch
import torch.nn as nn
import torch.nn.functional as F


@dataclass
class Config:
    vocab_size: int = 32000
    n_layer: int = 12
    n_head: int = 12
    n_kv_head: int = 4
    d_model: int = 768
    d_ff: int = 2048
    seq_len: int = 2048
    rope_theta: float = 10000.0
    ternary: bool = False
    # multiplier-free experiments (research/MULTIPLIER_FREE.md)
    wgrad_ternary: bool = False   # weight gradient from activations rounded to {-a, 0, +a}
    attn_lns: float = 0.0         # >1: attention products on signed powers of this base
    lin8: str = ""                # 8-bit linear operands: "fmt" or "fwd_fmt:bwd_fmt"
    lin8_scale: bool = False      # per-tensor amax scaling before rounding


# Ladder sizes: total parameters with vocab 32000 and tied embeddings.
SIZES = {
    "13m": dict(n_layer=6, n_head=4, n_kv_head=2, d_model=256, d_ff=704),
    "25m": dict(n_layer=8, n_head=6, n_kv_head=2, d_model=384, d_ff=1024),
    "52m": dict(n_layer=12, n_head=8, n_kv_head=4, d_model=512, d_ff=1408),
    "100m": dict(n_layer=12, n_head=12, n_kv_head=4, d_model=768, d_ff=2048),
    # TRI CLAW: ternary blocks sized for the XC7A200T's block RAM (6.45M weights,
    # 77% of 13.46 Mb at 1.6 bits); meant for the 8K vocabulary (data8k)
    "tc": dict(n_layer=6, n_head=5, n_kv_head=1, d_model=320, d_ff=864),
}


def rms_norm(x, eps=1e-6):
    return F.rms_norm(x, (x.size(-1),), eps=eps)


class _TernaryWgrad(torch.autograd.Function):
    """y = x_q W^T as usual; in the backward pass the weight gradient is taken
    against activations rounded to three levels per token, so dW = g^T x_t is a
    signed sum of gradient rows (the per-token scale is O(d) work, not O(d^2))."""

    @staticmethod
    def forward(ctx, xq, wq, xt):
        ctx.save_for_backward(wq, xt)
        ctx.dtypes = (xq.dtype, wq.dtype)
        return F.linear(xq, wq)

    @staticmethod
    def backward(ctx, gy):
        wq, xt = ctx.saved_tensors
        gx = (gy @ wq.to(gy.dtype)).to(ctx.dtypes[0])
        gw = gy.reshape(-1, gy.size(-1)).t() @ xt.reshape(-1, xt.size(-1)).to(gy.dtype)
        return gx, gw.to(ctx.dtypes[1]), None


class BitLinear(nn.Linear):
    """BitNet b1.58 linear layer (Ma et al. 2024, training-tips reference code)."""
    wgrad_ternary = False

    def forward(self, x):
        x = rms_norm(x)
        s = 127.0 / x.abs().amax(dim=-1, keepdim=True).clamp(min=1e-5)
        xq = x + ((x * s).round().clamp(-128, 127) / s - x).detach()
        w = self.weight
        g = w.abs().mean().clamp(min=1e-5)
        w = w + ((w / g).round().clamp(-1, 1) * g - w).detach()
        if self.wgrad_ternary:
            xd = x.detach()
            a = xd.abs().mean(dim=-1, keepdim=True).clamp(min=1e-5)
            return _TernaryWgrad.apply(xq, w, (xd / a).round().clamp(-1, 1) * a)
        return F.linear(xq, w)


def _fp_magnitudes(e_bits, m_bits, bias, top_nan):
    """Non-negative values of a small IEEE-style float (subnormals included)."""
    vals = {0.0}
    for m in range(1, 1 << m_bits):
        vals.add(m / (1 << m_bits) * 2.0 ** (1 - bias))
    for e in range(1, (1 << e_bits) - (0 if top_nan else 1)):
        for m in range(1 << m_bits):
            if top_nan and e == (1 << e_bits) - 1 and m == (1 << m_bits) - 1:
                continue  # e4m3: the single NaN code
            vals.add((1 + m / (1 << m_bits)) * 2.0 ** (e - bias))
    return sorted(vals)


def _tnf_magnitudes(trits, m_bits):
    """Ternary-exponent float packed into 8 binary bits (trinity-fpga
    research/block/four_families.py levels()): e in [-(3^t)//2, (3^t)//2]."""
    half = 3 ** trits // 2
    return sorted({0.0} | {(1 + m / (1 << m_bits)) * 2.0 ** e
                           for e in range(-half, half + 1) for m in range(1 << m_bits)})


FORMATS8 = {
    "e4m3": _fp_magnitudes(4, 3, 7, top_nan=True),        # max 448, OCP FP8
    "e5m2": _fp_magnitudes(5, 2, 15, top_nan=False),      # max 57344, OCP FP8
    "tnf8e2": _tnf_magnitudes(2, 3),                      # 73 values, TNF8 of the paper
    "tnf8e3": _tnf_magnitudes(3, 2),                      # 109 values, GF-T8 / TNF8 at Et=3
}
_TABLES = {}


def _table(name, device):
    key = (name, str(device))
    if key not in _TABLES:
        v = torch.tensor(FORMATS8[name], dtype=torch.float32, device=device)
        _TABLES[key] = (v, (v[1:] + v[:-1]) / 2)
    return _TABLES[key]


def round8(x, name, scale):
    """Round to the nearest value of an 8-bit format (saturating); optionally after
    per-tensor scaling so the tensor's max magnitude maps to the format's max."""
    v, mid = _table(name, x.device)
    xf = x.float()
    s = (v[-1] / xf.abs().amax().clamp(min=1e-30)) if scale else None
    y = xf * s if scale else xf
    q = torch.sign(y) * v[torch.bucketize(y.abs(), mid)]
    return (q / s if scale else q).to(x.dtype)


class _Fake8(torch.autograd.Function):
    """Linear layer whose three matmuls take 8-bit operands: x and W forward,
    the output gradient backward (fp32 master weights stay outside)."""

    @staticmethod
    def forward(ctx, x, w, fwd, bwd, scale):
        xq, wq = round8(x, fwd, scale), round8(w, fwd, scale)
        ctx.save_for_backward(xq, wq)
        ctx.bwd, ctx.scale, ctx.dtypes = bwd, scale, (x.dtype, w.dtype)
        return F.linear(xq, wq)

    @staticmethod
    def backward(ctx, gy):
        xq, wq = ctx.saved_tensors
        g = round8(gy, ctx.bwd, ctx.scale)
        gx = (g @ wq.to(g.dtype)).to(ctx.dtypes[0])
        gw = g.reshape(-1, g.size(-1)).t() @ xq.reshape(-1, xq.size(-1)).to(g.dtype)
        return gx, gw.to(ctx.dtypes[1]), None, None, None


class Fake8Linear(nn.Linear):
    fmt_fwd, fmt_bwd, scale = "e4m3", "e5m2", False

    def forward(self, x):
        return _Fake8.apply(x, self.weight, self.fmt_fwd, self.fmt_bwd, self.scale)


def lns_round(x, base):
    """Nearest signed power of `base`, rounded in the log domain, as an 8-bit
    code would hold it (sign + 7-bit exponent in [-96, 31]); tiny values -> 0.
    base == 1 returns x unchanged (used to test the attention path)."""
    if base == 1.0:
        return x
    mag = x.abs().float()
    e = torch.round(torch.log(mag.clamp(min=1e-38)) / math.log(base))
    q = torch.sign(x.float()) * torch.pow(base, e.clamp(-96, 31))
    return torch.where(e < -96, torch.zeros_like(q), q).to(x.dtype)


def _lns_q(x, inv_lb, lb):
    """lns_round for float32 tensors with precomputed 1/log(base), log(base)."""
    e = torch.round(torch.log(x.abs().clamp(min=1e-38)) * inv_lb)
    q = torch.sign(x) * torch.exp(e.clamp(-96, 31) * lb)
    return torch.where(e < -96, torch.zeros_like(q), q)


def _causal_mask(T, device):
    return torch.ones(T, T, dtype=torch.bool, device=device).triu(1)


def _lns_attn_fwd(q, k, v, inv_lb, lb, scale):
    qs, ks, vs = (_lns_q(t.float(), inv_lb, lb) for t in (q, k, v))
    s = (qs @ ks.transpose(-1, -2)) * scale
    p = torch.softmax(s.masked_fill(_causal_mask(s.size(-1), s.device), float("-inf")), dim=-1)
    return _lns_q(p, inv_lb, lb) @ vs, qs, ks, vs


def _lns_attn_bwd(qs, ks, vs, go, inv_lb, lb, scale):
    s = (qs @ ks.transpose(-1, -2)) * scale
    p = torch.softmax(s.masked_fill(_causal_mask(s.size(-1), s.device), float("-inf")), dim=-1)
    ps, gos = _lns_q(p, inv_lb, lb), _lns_q(go.float(), inv_lb, lb)
    gv = ps.transpose(-1, -2) @ gos
    gp = gos @ vs.transpose(-1, -2)
    gs = _lns_q(p * (gp - (gp * p).sum(-1, keepdim=True)) * scale, inv_lb, lb)
    return gs @ ks, gs.transpose(-1, -2) @ qs, gv


_lns_attn_fwd_c = torch.compile(_lns_attn_fwd, dynamic=False)
_lns_attn_bwd_c = torch.compile(_lns_attn_bwd, dynamic=False)


class _LNSAttention(torch.autograd.Function):
    """Causal attention whose six matrix products (two forward, four backward)
    take operands rounded to signed powers of `base`, as a logarithmic-number
    datapath would (a product there is an exponent addition). Accumulation and
    softmax stay exact, i.e. an ideal wide accumulator; P is recomputed.
    base == 1 skips the rounding (used to test the attention path)."""

    @staticmethod
    def forward(ctx, q, k, v, base):
        scale = q.size(-1) ** -0.5
        lb = math.log(base) if base != 1.0 else 0.0
        inv = 1.0 / lb if lb else 0.0
        fwd = _lns_attn_fwd_c if q.is_cuda else _lns_attn_fwd
        if base == 1.0:
            qs, ks, vs = q.float(), k.float(), v.float()
            p = _causal_probs(qs, ks, scale)
            out = p @ vs
        else:
            out, qs, ks, vs = fwd(q, k, v, inv, lb, scale)
        ctx.save_for_backward(qs, ks, vs)
        ctx.base, ctx.scale, ctx.lb, ctx.inv, ctx.dtype = base, scale, lb, inv, q.dtype
        return out.to(q.dtype)

    @staticmethod
    def backward(ctx, go):
        qs, ks, vs = ctx.saved_tensors
        if ctx.base == 1.0:
            p = _causal_probs(qs, ks, ctx.scale)
            g = go.float()
            gv = p.transpose(-1, -2) @ g
            gp = g @ vs.transpose(-1, -2)
            gs = p * (gp - (gp * p).sum(-1, keepdim=True)) * ctx.scale
            gq, gk = gs @ ks, gs.transpose(-1, -2) @ qs
        else:
            bwd = _lns_attn_bwd_c if go.is_cuda else _lns_attn_bwd
            gq, gk, gv = bwd(qs, ks, vs, go, ctx.inv, ctx.lb, ctx.scale)
        return gq.to(ctx.dtype), gk.to(ctx.dtype), gv.to(ctx.dtype), None


def _causal_probs(q, k, scale):
    s = (q.float() @ k.float().transpose(-1, -2)) * scale
    return torch.softmax(s.masked_fill(_causal_mask(s.size(-1), s.device), float("-inf")), dim=-1)


@torch.compiler.disable
def lns_attention(q, k, v, base):
    """Opaque to the model-level compile; the helpers inside are compiled on their own."""
    return _LNSAttention.apply(q, k, v, base)


def rope_cache(seq_len, head_dim, theta, device):
    inv = 1.0 / theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)
    t = torch.arange(seq_len, device=device).float()
    freqs = torch.outer(t, inv)
    return freqs.cos(), freqs.sin()


def apply_rope(x, cos, sin):
    # x: (B, T, H, D); cos/sin: (T, D/2)
    x1, x2 = x.float().chunk(2, dim=-1)
    cos, sin = cos[None, :, None, :], sin[None, :, None, :]
    return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).type_as(x)


class Attention(nn.Module):
    def __init__(self, cfg, Lin):
        super().__init__()
        self.lns = cfg.attn_lns
        self.h, self.kv = cfg.n_head, cfg.n_kv_head
        self.hd = cfg.d_model // cfg.n_head
        d = cfg.d_model
        self.q = Lin(d, self.h * self.hd, bias=False)
        self.k = Lin(d, self.kv * self.hd, bias=False)
        self.v = Lin(d, self.kv * self.hd, bias=False)
        self.o = Lin(self.h * self.hd, d, bias=False)

    def forward(self, x, cos, sin, cache=None):
        B, T, _ = x.shape
        q = apply_rope(self.q(x).view(B, T, self.h, self.hd), cos, sin).transpose(1, 2)
        k = apply_rope(self.k(x).view(B, T, self.kv, self.hd), cos, sin).transpose(1, 2)
        v = self.v(x).view(B, T, self.kv, self.hd).transpose(1, 2)
        if cache is not None:
            k, v = cache.update(k, v)
        # A single-token decode step attends to the whole cache; everything else is causal.
        causal = T > 1
        if self.kv != self.h:
            k = k.repeat_interleave(self.h // self.kv, dim=1)
            v = v.repeat_interleave(self.h // self.kv, dim=1)
        if self.lns > 1 and cache is None:
            y = lns_attention(q, k, v, self.lns)
        else:
            y = F.scaled_dot_product_attention(q, k, v, is_causal=causal)
        return self.o(y.transpose(1, 2).reshape(B, T, self.h * self.hd))


class MLP(nn.Module):
    def __init__(self, cfg, Lin):
        super().__init__()
        self.gate = Lin(cfg.d_model, cfg.d_ff, bias=False)
        self.up = Lin(cfg.d_model, cfg.d_ff, bias=False)
        self.down = Lin(cfg.d_ff, cfg.d_model, bias=False)

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


class Block(nn.Module):
    def __init__(self, cfg, Lin):
        super().__init__()
        self.n1 = nn.RMSNorm(cfg.d_model, eps=1e-6)
        self.attn = Attention(cfg, Lin)
        self.n2 = nn.RMSNorm(cfg.d_model, eps=1e-6)
        self.mlp = MLP(cfg, Lin)

    def forward(self, x, cos, sin, cache=None):
        x = x + self.attn(self.n1(x), cos, sin, cache)
        return x + self.mlp(self.n2(x))


class KVCache:
    """Per-layer key/value buffers for batched generation of equal-length prompts."""

    def __init__(self, B, kv_heads, head_dim, max_len, device, dtype):
        self.k = torch.zeros(B, kv_heads, max_len, head_dim, device=device, dtype=dtype)
        self.v = torch.zeros_like(self.k)
        self.pos = 0

    def update(self, k, v):
        T = k.size(2)
        self.k[:, :, self.pos:self.pos + T] = k
        self.v[:, :, self.pos:self.pos + T] = v
        self.pos += T
        return self.k[:, :, :self.pos], self.v[:, :, :self.pos]


class GPT(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg
        Lin = BitLinear if cfg.ternary else Fake8Linear if cfg.lin8 else nn.Linear
        self.wte = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.blocks = nn.ModuleList(Block(cfg, Lin) for _ in range(cfg.n_layer))
        self.norm = nn.RMSNorm(cfg.d_model, eps=1e-6)
        cos, sin = rope_cache(cfg.seq_len * 2, cfg.d_model // cfg.n_head, cfg.rope_theta, "cpu")
        self.register_buffer("cos", cos, persistent=False)
        self.register_buffer("sin", sin, persistent=False)
        self.apply(self._init)
        fwd8, _, bwd8 = cfg.lin8.partition(":")
        for m in self.modules():
            if isinstance(m, BitLinear):
                m.wgrad_ternary = cfg.wgrad_ternary
            if isinstance(m, Fake8Linear):
                m.fmt_fwd, m.fmt_bwd, m.scale = fwd8, bwd8 or fwd8, cfg.lin8_scale
        for name, p in self.named_parameters():
            if name.endswith(("attn.o.weight", "mlp.down.weight")):
                nn.init.normal_(p, std=0.02 / math.sqrt(2 * cfg.n_layer))

    @staticmethod
    def _init(m):
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, std=0.02)

    def forward(self, idx, targets=None, caches=None):
        T = idx.size(1)
        start = caches[0].pos if caches else 0
        cos, sin = self.cos[start:start + T], self.sin[start:start + T]
        x = self.wte(idx)
        for i, blk in enumerate(self.blocks):
            x = blk(x, cos, sin, caches[i] if caches else None)
        logits = F.linear(self.norm(x), self.wte.weight)
        if targets is None:
            return logits
        loss = F.cross_entropy(logits.float().view(-1, logits.size(-1)), targets.reshape(-1),
                               reduction="none")
        return loss.view(targets.shape)

    def num_params(self):
        return sum(p.numel() for p in self.parameters())

    def flops_per_token(self, seq_len=None):
        """Training FLOPs per token: 6 * matmul params + causal attention (Kaplan 2020).

        Every parameter except the norms takes part in a matmul (the tied
        embedding matrix does so as the output head)."""
        c = self.cfg
        norms = sum(p.numel() for m in self.modules() if isinstance(m, nn.RMSNorm)
                    for p in m.parameters())
        return 6 * (self.num_params() - norms) + 6 * c.n_layer * (seq_len or c.seq_len) * c.d_model

    @torch.no_grad()
    def generate(self, prompt, max_new, temperature=0.2, top_p=0.95, cache_dtype=torch.float32,
                 stop_fn=None, check_every=16):
        """Sample continuations for a batch of equal-length prompts (B, T).

        stop_fn(tokens) is called every `check_every` steps with the tokens
        generated so far; returning True ends generation early."""
        B, T = prompt.shape
        c = self.cfg
        caches = [KVCache(B, c.n_kv_head, c.d_model // c.n_head, T + max_new, prompt.device,
                          cache_dtype) for _ in range(c.n_layer)]
        logits = self(prompt, caches=caches)[:, -1]
        out = []
        for i in range(max_new):
            if temperature > 0:
                probs = torch.softmax(logits.float() / temperature, dim=-1)
                sp, si = probs.sort(dim=-1, descending=True)
                sp[sp.cumsum(-1) - sp > top_p] = 0
                nxt = si.gather(-1, torch.multinomial(sp / sp.sum(-1, keepdim=True), 1))
            else:
                nxt = logits.argmax(-1, keepdim=True)
            out.append(nxt)
            if stop_fn and (i + 1) % check_every == 0 and stop_fn(torch.cat(out, dim=1)):
                break
            logits = self(nxt, caches=caches)[:, -1]
        return torch.cat(out, dim=1)


def build(size, ternary=False, vocab_size=32000, seq_len=2048, **extra):
    return GPT(Config(vocab_size=vocab_size, seq_len=seq_len, ternary=ternary, **SIZES[size],
                      **extra))


def config_dict(model):
    return asdict(model.cfg)