"""OLMo2-style decoder LM with the ZBP physical/digital partition (part-1 'mixscoreqk,ffn' convention):
physical = score-space attention core (per head x row) + whole SwiGLU FFN branch (per token);
digital = embeddings, RMSNorms outside blocks, q/k/v/proj linears, score product, readout head."""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

from .zbp import ZBPBlock, ZBPConfig


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

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


class _Core(nn.Module):
    """softmax(QK^T/sqrt(dh)) V from concatenated qkv [B,T,3d] -> [B,T,d]; explicit (forward-AD friendly)."""

    def __init__(self, d, heads):
        super().__init__()
        self.d, self.h = d, heads

    def forward(self, qkv):
        B, T, _ = qkv.shape
        q, k, v = qkv.split(self.d, dim=-1)
        dh = self.d // self.h
        q = q.view(B, T, self.h, dh).transpose(1, 2)
        k = k.view(B, T, self.h, dh).transpose(1, 2)
        v = v.view(B, T, self.h, dh).transpose(1, 2)
        att = (q @ k.transpose(-2, -1)) / math.sqrt(dh)
        mask = torch.ones(T, T, dtype=torch.bool, device=qkv.device).tril()
        att = att.masked_fill(~mask, float("-inf")).softmax(-1)
        return (att @ v).transpose(1, 2).reshape(B, T, self.d)


class _SwiGLU(nn.Module):
    """RMS -> (W1 x) * SiLU(W3 x) -> W2, one per-token physical branch (identity skip outside)."""

    def __init__(self, d, hidden, scale):
        super().__init__()
        self.norm = RMSNorm(d)
        self.w1 = nn.Linear(d, hidden, bias=False)
        self.w3 = nn.Linear(d, hidden, bias=False)
        self.w2 = nn.Linear(hidden, d, bias=False)
        with torch.no_grad():
            self.w2.weight.mul_(scale)

    def forward(self, x):
        h = self.norm(x)
        return self.w2(self.w1(h) * F.silu(self.w3(h)))


class Block(nn.Module):
    def __init__(self, d, heads, hidden, cfg, bp, scale, li):
        super().__init__()
        self.norm_a = RMSNorm(d)
        self.qkv = nn.Linear(d, 3 * d, bias=False)
        core = ZBPBlock(_Core(d, heads), cfg, batch_dims=1, name=f"L{li}.core")
        core.score_probe = {"d": d, "heads": heads, "dk": d, "dv": d, "window": None, "qk": True}
        self.core = core
        self.proj = nn.Linear(d, d, bias=False)
        with torch.no_grad():
            self.proj.weight.mul_(scale)
        self.ffn = ZBPBlock(_SwiGLU(d, hidden, scale), cfg, skip="identity", batch_dims=2, name=f"L{li}.ffn")

    def forward(self, x):
        x = x + self.proj(self.core(self.qkv(self.norm_a(x))))
        return self.ffn(x)


class ScalingLM(nn.Module):
    def __init__(self, vocab, d, layers, heads, seq_len, ffn_mult="8/3", cfg=None):
        super().__init__()
        cfg = cfg or ZBPConfig(mode="bp")
        bp = cfg.replace(mode="bp")
        mult = eval(str(ffn_mult)) if isinstance(ffn_mult, str) else ffn_mult
        hidden = int(round(mult * d / 64) * 64)
        scale = 1 / math.sqrt(2 * layers)
        self.tok = nn.Embedding(vocab, d)
        self.pos = nn.Embedding(seq_len, d)
        nn.init.normal_(self.tok.weight, std=0.02); nn.init.normal_(self.pos.weight, std=0.02)
        self.blocks = nn.Sequential(*[Block(d, heads, hidden, cfg, bp, scale, i) for i in range(layers)])
        self.norm_f = RMSNorm(d)
        self.head = nn.Linear(d, vocab, bias=False)
        self.seq_len = seq_len

    def forward(self, idx):
        B, T = idx.shape
        x = self.tok(idx) + self.pos(torch.arange(T, device=idx.device))[None]
        return self.head(self.norm_f(self.blocks(x)))