"""TENP expert-neuron projection plus optional energy control. No OLS."""

from __future__ import annotations

import torch

PROJ_EPS = 1e-6


def gated_signed_projection(
    z: torch.Tensor,
    y: torch.Tensor,
    gate: torch.Tensor,
    W_down: torch.Tensor,
    eps: float = PROJ_EPS,
) -> torch.Tensor:
    """Per-token score `q = g * p` with signed projection onto expert output y.

    z: [n, d_int] SwiGLU mid activations
    y: [n, d] expert output (not yet multiplied by router gate)
    gate: [n] router weight (top-1 after renormalize is ~1)
    W_down: [d, d_int] down_proj.weight
    returns q: [n, d_int]
    """
    w = W_down.detach().to(device=y.device, dtype=torch.float32)
    y_f = y.float()
    z_f = z.float()
    g = gate.float().reshape(-1, 1)
    y_norm = y_f.norm(dim=-1, keepdim=True).clamp_min(eps)
    p = z_f * (y_f @ w) / y_norm
    return p * g


def gated_energy(z: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
    """Control score `g^2 z^2` (top-1 gate is ~1, so this is mostly z^2)."""
    g = gate.float().reshape(-1, 1)
    return (g * z.float()).square()