"""MN-GLM — a GLM-5.2-class decoder skeleton on the counter-synapse method.

Modern attention stack (RMSNorm + GQA + RoPE + optional QK-norm) with a sparse Counter-MoE FFN
(plan M4, grouped/stacked kernels). Every linear is a counter layer (ternary + 6-bit counter,
self-updating), so the GLM body inherits the method's sub-byte weights + in-state optimizer.

What this adds over the repo's GPT (the GLM build-list): RMSNorm (was LayerNorm), GQA (was full
MHA), RoPE (was learned positional embeddings), QK-norm, and SwiGLU experts (default; gate/up/down,
the GLM/Llama expert MLP) in the sparse Counter-MoE FFN -- all on the grouped/stacked kernels.
Reversible wrapping (the activation-memory lever) composes on top: `ReversibleMNGLM` gives the same
GLM stack at O(1)-in-depth activation memory.

Pure PyTorch, CPU/CUDA.
"""
from __future__ import annotations

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

from .baselines import make_linear
from .counter import CompactCounterLinear
from .moe_ffn import CounterMoEFFN
from .reversible import ReversibleCouplingBlock, ReversibleSequence

__all__ = ["RMSNorm", "GLMAttention", "GLMBlock", "MNGLM", "ReversibleMNGLM"]


class RMSNorm(nn.Module):
    """Root-mean-square layer norm (GLM/Llama style): no mean-subtraction, one scale per channel."""

    def __init__(self, d: int, eps: float = 1e-5) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d))
        self.eps = eps

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


def _rope_cache(T: int, head_dim: int, device, dtype, base: float = 10000.0):
    """cos/sin tables [T, head_dim] for rotary position embedding (NeoX/Llama 'rotate_half' form)."""
    inv = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    freqs = torch.outer(torch.arange(T, device=device).float(), inv)   # [T, head_dim/2]
    emb = torch.cat([freqs, freqs], dim=-1)                             # [T, head_dim]
    return emb.cos().to(dtype), emb.sin().to(dtype)


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)


def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    # x [B, H, T, head_dim]; cos/sin [T, head_dim]
    return x * cos[None, None] + _rotate_half(x) * sin[None, None]


class GLMAttention(nn.Module):
    """Grouped-query attention with RoPE and optional QK-norm; all projections are counter linears.

    n_head query heads, n_kv_head key/value heads (GQA: n_head % n_kv_head == 0). The KV heads are
    repeated to match the query heads before SDPA, shrinking the KV projection / cache by the ratio.
    """

    def __init__(self, d: int, n_head: int, n_kv_head: int, kind: str, ckw: dict,
                 qk_norm: bool = True) -> None:
        super().__init__()
        if d % n_head != 0:
            raise ValueError("n_embd must be divisible by n_head")
        if n_head % n_kv_head != 0:
            raise ValueError("n_head must be divisible by n_kv_head (GQA)")
        self.nh, self.nkv = n_head, n_kv_head
        self.hd = d // n_head
        if self.hd % 2 != 0:
            raise ValueError("head_dim must be even for RoPE")
        self.d = d
        self.q = make_linear(kind, d, n_head * self.hd, 1.0, **ckw)
        self.k = make_linear(kind, d, n_kv_head * self.hd, 1.0, **ckw)
        self.v = make_linear(kind, d, n_kv_head * self.hd, 1.0, **ckw)
        self.o = make_linear(kind, d, d, 1.0, **ckw)
        self.qn = RMSNorm(self.hd) if qk_norm else None
        self.kn = RMSNorm(self.hd) if qk_norm else None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, _ = x.shape
        cos, sin = _rope_cache(T, self.hd, x.device, x.dtype)          # computed internally (from T)
        q = self.q(x).view(B, T, self.nh, self.hd).transpose(1, 2)     # [B, nh, T, hd]
        k = self.k(x).view(B, T, self.nkv, self.hd).transpose(1, 2)    # [B, nkv, T, hd]
        v = self.v(x).view(B, T, self.nkv, self.hd).transpose(1, 2)
        if self.qn is not None:
            q = self.qn(q); k = self.kn(k)                             # QK-norm (per-head-dim RMS)
        q = _apply_rope(q, cos, sin); k = _apply_rope(k, cos, sin)
        rep = self.nh // self.nkv
        if rep > 1:
            k = k.repeat_interleave(rep, dim=1)                        # GQA: broadcast KV to all q heads
            v = v.repeat_interleave(rep, dim=1)
        a = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        a = a.transpose(1, 2).reshape(B, T, self.d)
        return self.o(a)


def _counter_numeric(kw: dict) -> dict:
    return {k: kw[k] for k in ("C", "lr", "lr_scale") if k in kw}


def _glm_head_loss(model, h, targets, loss_chunk):
    """Shared head/loss for the GLM models: tied head + CE + MoE load-balance aux. loss_chunk>0
    chunks the head over rows so the [B,T,V] logits tensor (V~50k) is never materialized -- returns
    (None, loss). loss_chunk=0 returns (logits, loss)."""
    if targets is not None and loss_chunk > 0:
        d = model.d
        hf = h.reshape(-1, d); tf = targets.reshape(-1)
        total = hf.new_zeros(()); n = hf.shape[0]
        for i in range(0, n, loss_chunk):
            lg = model.head(hf[i:i + loss_chunk])
            total = total + F.cross_entropy(lg, tf[i:i + loss_chunk], reduction="sum")
        loss = total / n
        aux = model.aux_loss()
        if aux is not None:
            loss = loss + aux
        return None, loss
    logits = model.head(h)
    loss = None
    if targets is not None:
        loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))
        aux = model.aux_loss()
        if aux is not None:
            loss = loss + aux
    return logits, loss


class GLMBlock(nn.Module):
    """RMSNorm-pre-norm GLM block: attention sublayer + Counter-MoE FFN sublayer, both residual."""

    def __init__(self, d: int, n_head: int, n_kv_head: int, kind: str, ckw: dict,
                 n_experts: int, top_k: int, qk_norm: bool, grouped: bool,
                 aux_loss_weight: float = 1e-2, swiglu: bool = True,
                 moe_dtype: str = "fp32") -> None:
        super().__init__()
        self.n_embd = d
        self.n1 = RMSNorm(d)
        self.attn = GLMAttention(d, n_head, n_kv_head, kind, ckw, qk_norm)
        self.n2 = RMSNorm(d)
        self.ffn = CounterMoEFFN(d, n_experts=n_experts, top_k=top_k, grouped=grouped,
                                 aux_loss_weight=aux_loss_weight, swiglu=swiglu,
                                 compute_dtype=moe_dtype, **_counter_numeric(ckw))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.n1(x))
        x = x + self.ffn(self.n2(x))
        return x


class MNGLM(nn.Module):
    """GLM-5.2-class decoder on the counter method: RMSNorm + GQA + RoPE + Counter-MoE, tied head."""

    def __init__(self, vocab_size: int, n_embd: int, n_layer: int, n_head: int, n_kv_head: int,
                 block_size: int, *, kind: str = "counter_packed", n_experts: int = 8, top_k: int = 2,
                 qk_norm: bool = True, grouped: bool = True, aux_loss_weight: float = 1e-2,
                 swiglu: bool = True, moe_dtype: str = "fp32", **counter_kw) -> None:
        super().__init__()
        self.block_size = block_size
        self.d = n_embd
        self.n_head = n_head
        self.tok = nn.Embedding(vocab_size, n_embd)
        nn.init.normal_(self.tok.weight, std=0.02)
        self.blocks = nn.ModuleList(
            GLMBlock(n_embd, n_head, n_kv_head, kind, counter_kw, n_experts, top_k,
                     qk_norm, grouped, aux_loss_weight, swiglu, moe_dtype)
            for _ in range(n_layer))
        self.nf = RMSNorm(n_embd)
        self.head = nn.Linear(n_embd, vocab_size, bias=False)
        self.head.weight = self.tok.weight                             # tie

    def forward(self, idx: torch.Tensor, targets: torch.Tensor | None = None, loss_chunk: int = 0):
        B, T = idx.shape
        if T > self.block_size:
            raise ValueError(f"sequence length {T} exceeds block size {self.block_size}")
        x = self.tok(idx)
        for b in self.blocks:
            x = b(x)
        h = self.nf(x)
        return _glm_head_loss(self, h, targets, loss_chunk)

    def aux_loss(self) -> torch.Tensor | None:
        total = None
        for m in self.modules():
            w = float(getattr(m, "aux_loss_weight", 0.0) or 0.0)
            if w > 0.0 and hasattr(m, "last_aux_loss"):
                term = w * m.last_aux_loss
                total = term if total is None else total + term
        return total

    def counter_layers(self) -> list[CompactCounterLinear]:
        return [m for m in self.modules() if isinstance(m, CompactCounterLinear)]

    def total_counter_coeffs(self) -> int:
        """ALL counter weights: attention counter linears + the MoE experts (which live in stacked
        buffers, not CompactCounterLinear instances, so counter_layers() alone undercounts)."""
        n = sum(c.in_features * c.out_features for c in self.counter_layers())
        for mod in self.modules():
            if isinstance(mod, CounterMoEFFN):
                n += mod.coeff_count()
        return n

    def trainable_parameters(self):
        return [p for p in self.parameters() if p.requires_grad]


class _GLMAttnSub(nn.Module):
    """F for the reversible coupling: RMSNorm -> GQA(+RoPE, QK-norm) -> out-proj, maps [.,d]->[.,d]."""

    def __init__(self, d, n_head, n_kv_head, kind, ckw, qk_norm):
        super().__init__()
        self.d = d
        self.n = RMSNorm(d)
        self.attn = GLMAttention(d, n_head, n_kv_head, kind, ckw, qk_norm)

    def forward(self, x):
        return self.attn(self.n(x))


class _GLMFFNSub(nn.Module):
    """G for the reversible coupling: RMSNorm -> Counter-MoE FFN, maps [.,d]->[.,d]."""

    def __init__(self, d, kind, ckw, n_experts, top_k, grouped, aux_loss_weight, swiglu,
                 moe_dtype="fp32"):
        super().__init__()
        self.d = d
        self.n = RMSNorm(d)
        self.ffn = CounterMoEFFN(d, n_experts=n_experts, top_k=top_k, grouped=grouped,
                                 aux_loss_weight=aux_loss_weight, swiglu=swiglu,
                                 compute_dtype=moe_dtype, **_counter_numeric(ckw))

    def forward(self, x):
        return self.ffn(self.n(x))


class ReversibleMNGLM(nn.Module):
    """MN-GLM with O(1)-in-depth activation memory: the GLM blocks are wrapped in a reversible
    coupling stack (F = attention sublayer, G = MoE FFN sublayer). Same GLM-5.2 components + tuned
    counter knobs as MNGLM; adds the activation-memory lever (anchor_every is the speed/memory knob,
    2-4 recommended per PERF_ANATOMY). The two reversible streams are the duplicated embedding."""

    def __init__(self, vocab_size: int, n_embd: int, n_layer: int, n_head: int, n_kv_head: int,
                 block_size: int, *, kind: str = "counter_packed", n_experts: int = 8, top_k: int = 2,
                 qk_norm: bool = True, grouped: bool = True, aux_loss_weight: float = 1e-2,
                 anchor_every: int = 0, swiglu: bool = True, moe_dtype: str = "fp32",
                 **counter_kw) -> None:
        super().__init__()
        self.block_size = block_size
        self.d = n_embd
        self.tok = nn.Embedding(vocab_size, n_embd)
        nn.init.normal_(self.tok.weight, std=0.02)
        blocks = [ReversibleCouplingBlock(
                      2 * n_embd,
                      F=_GLMAttnSub(n_embd, n_head, n_kv_head, kind, counter_kw, qk_norm),
                      G=_GLMFFNSub(n_embd, kind, counter_kw, n_experts, top_k, grouped,
                                   aux_loss_weight, swiglu, moe_dtype))
                  for _ in range(n_layer)]
        self.rev = ReversibleSequence(blocks, anchor_every=anchor_every)
        self.nf = RMSNorm(n_embd)
        self.head = nn.Linear(n_embd, vocab_size, bias=False)
        self.head.weight = self.tok.weight

    def forward(self, idx: torch.Tensor, targets: torch.Tensor | None = None, loss_chunk: int = 0):
        B, T = idx.shape
        if T > self.block_size:
            raise ValueError(f"sequence length {T} exceeds block size {self.block_size}")
        e = self.tok(idx)
        x = torch.cat([e, e], dim=-1)                                  # two reversible streams
        x = self.rev(x)
        h = self.nf(x[..., :self.d] + x[..., self.d:])                 # recombine
        return _glm_head_loss(self, h, targets, loss_chunk)

    def aux_loss(self) -> torch.Tensor | None:
        total = None
        for m in self.modules():
            w = float(getattr(m, "aux_loss_weight", 0.0) or 0.0)
            if w > 0.0 and hasattr(m, "last_aux_loss"):
                term = w * m.last_aux_loss
                total = term if total is None else total + term
        return total

    def counter_layers(self) -> list[CompactCounterLinear]:
        return [m for m in self.modules() if isinstance(m, CompactCounterLinear)]

    def total_counter_coeffs(self) -> int:
        """ALL counter weights: attention counter linears + the MoE experts (which live in stacked
        buffers, not CompactCounterLinear instances, so counter_layers() alone undercounts)."""
        n = sum(c.in_features * c.out_features for c in self.counter_layers())
        for mod in self.modules():
            if isinstance(mod, CounterMoEFFN):
                n += mod.coeff_count()
        return n

    def trainable_parameters(self):
        return [p for p in self.parameters() if p.requires_grad]