"""GemmaHyperExpert -- the per-token GENERATED low-rank student that replaces a
Gemma-4 decoder layer's whole feed-forward/MoE BLOCK.

Trained LAYER-LOCALLY by feature distillation against the teacher's cached block
I/O (X = block input = pre_feedforward_layernorm input ; Y = block output = X+D),
with NO access to the teacher's weights at train time -- only the cached (X, Y)
activations.  This is the Gemma adaptation of the StarCoder testbed's HyperExpert
(../hyper_expert.py), with two deliberate changes:

  1. RESIDUAL parameterization.  The teacher block is a residual: Y = X + D.  So
     the student models the DELTA on top of the carried residual:
         Yhat = X + base(X) + V @ gelu(U @ X)
     and is trained to minimize relMSE(Yhat, Y).  Modeling D (not Y outright) is
     the easy/stable target and gives an exact identity warm start.

  2. FROM-SCRATCH init (no src_mlp).  We have no teacher MLP weights to subset, so
     base is a learned width-b FFN.  We ZERO the two output paths (b_proj and g_v)
     so at init  base(X)=0  and the generated correction=0  =>  Yhat = X exactly
     (the block is the identity residual).  Init relMSE is then ||D||^2/||Y||^2,
     a finite, meaningful starting point that training drives down.  enc, g_u and
     b_fc get small normal init; b_proj and g_v start at zero and pick up gradient
     immediately (b_proj from grad*b_fc_act, g_v from grad*h), recovering in a few
     steps -- the validated warm-start trick from the testbed.

Architecture (hidden d=2816 for Gemma's text decoder), the validated winning shape
"large hypernetwork, small generated rank":
    encoder   : z = gelu(Enc(x)),           Enc: Linear(d -> c)
    generator : U = (G_U z).view(r, d),     G_U: Linear(c -> r*d, no bias)
                V = (G_V z).view(d, r),      G_V: Linear(c -> d*r, no bias)
    base      : base(x) = b_proj(gelu(b_fc(x))),  widths d->b->d
    output    : Yhat = x + base(x) + V @ gelu(U @ x)

Per-layer params ~ d*c (enc) + 2*c*r*d (generators) + 2*d*b (base).  With d=2816,
c=9728, r=4, b=512 this is ~250M/layer => ~7.5B over 30 layers (teacher/3.5).

This is the PER-TOKEN core (chunk=1 in the testbed).  The cached activations are
stored flattened ([n_tokens, d], sequence/boundary structure not retained), so the
chunked / sentence-boundary DECODE-SPEED variants are an assembly-time concern and
are intentionally omitted here -- per-token distillation is the capability being
trained.  gelu uses the tanh approximation to match the testbed exactly.
"""
import torch, torch.nn as nn, torch.nn.functional as F


def expert_params(d, c, r, b):
    """Exact parameter count of one GemmaHyperExpert (for sizing to a target B)."""
    enc = d * c + c
    g_u = c * (r * d)
    g_v = c * (d * r)
    b_fc = d * b + b
    b_proj = b * d + d
    return enc + g_u + g_v + b_fc + b_proj


class GemmaHyperExpert(nn.Module):
    def __init__(self, d, c, r, b, dtype=torch.bfloat16):
        super().__init__()
        self.d, self.c, self.r, self.b = d, c, r, b
        self.dtype = dtype

        self.enc = nn.Linear(d, c)                  # x -> z
        self.g_u = nn.Linear(c, r * d, bias=False)  # z -> U (r,d)
        self.g_v = nn.Linear(c, d * r, bias=False)  # z -> V (d,r)
        self.b_fc = nn.Linear(d, b)                 # base up
        self.b_proj = nn.Linear(b, d)               # base down

        nn.init.normal_(self.enc.weight, std=0.02); nn.init.zeros_(self.enc.bias)
        nn.init.normal_(self.g_u.weight, std=0.02)
        nn.init.zeros_(self.g_v.weight)             # generated correction = 0 at init
        nn.init.normal_(self.b_fc.weight, std=0.02); nn.init.zeros_(self.b_fc.bias)
        nn.init.zeros_(self.b_proj.weight); nn.init.zeros_(self.b_proj.bias)  # base = 0 at init

        self.to(dtype)

    def base(self, x):
        return self.b_proj(F.gelu(self.b_fc(x), approximate="tanh"))

    def forward(self, x):
        """x: [..., d] block input X.  Returns Yhat = X + base(X) + lowrank(X)."""
        shape = x.shape
        x2 = x.reshape(-1, self.d).to(self.dtype)        # [N, d]
        N = x2.shape[0]
        z = F.gelu(self.enc(x2), approximate="tanh")     # [N, c]
        U = self.g_u(z).view(N, self.r, self.d)          # [N, r, d]
        V = self.g_v(z).view(N, self.d, self.r)          # [N, d, r]
        h = F.gelu(torch.bmm(U, x2.unsqueeze(-1)).squeeze(-1), approximate="tanh")  # [N, r]
        res = torch.bmm(V, h.unsqueeze(-1)).squeeze(-1)  # [N, d]
        y = x2 + self.base(x2) + res                     # residual block output
        return y.view(shape).to(x.dtype)

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


if __name__ == "__main__":
    # quick sizing table + forward/shape/warm-start sanity on CPU
    d = 2816
    for (c, r, b) in [(512, 4, 512), (9728, 4, 512), (9856, 4, 640), (9984, 4, 512)]:
        p = expert_params(d, c, r, b)
        print(f"c={c} r={r} b={b}: {p/1e6:.1f}M/layer -> {p*30/1e9:.2f}B over 30 layers")
    torch.manual_seed(0)
    e = GemmaHyperExpert(d, 512, 4, 512, dtype=torch.float32)
    x = torch.randn(3, 7, d)
    y = e(x)
    assert y.shape == x.shape
    # warm start: Yhat == X exactly at init (base=0, g_v=0)
    print("warm-start max|Yhat-X| =", (y - x).abs().max().item(), "(should be ~0)")
    print("footprint check:", e.footprint(), "==", expert_params(d, 512, 4, 512))