"""Gradient projection + delta_S grad utilities. Imported by smoke and train.

Shape conventions in the v_hack / delta_S plumbing (jaxtyping-annotated):
- `r` = per-module SVD rank (delta_S dimension; varies per Linear)
- `k` = number of v_hack directions kept per module (after top-k slice and
        global noise-floor filter at load time)
- `V` is `[k, r]`, rows orthonormal in R^r and oriented hack-ward
- `g = delta_S.grad` is `[r]`
- `c = V @ g` is `[k]`
"""
from __future__ import annotations

import torch
from jaxtyping import Float


def per_token_logps(logits: torch.Tensor, ids: torch.Tensor) -> torch.Tensor:
    """log p(ids | logits) gathered token-wise.

    Uses F.cross_entropy (fused softmax+gather) so we never materialise the
    full [B, L, V] fp32 softmax. On Qwen3.5-2B with V=152k, G=8, L≈1500 the
    fp32 vocab tensor was ~7 GB per forward -- the difference between OOM and
    fit on a 96 GB card when the autograd graph is alive.
    """
    B, L, V = logits.shape
    # CE's internal log_softmax accumulates in fp32 (stable) but returns input dtype.
    # The output [B*L] is small, so upcast it to fp32 for downstream PPO ratio math.
    return -torch.nn.functional.cross_entropy(
        logits.reshape(-1, V), ids.reshape(-1), reduction="none"
    ).float().view(B, L)


def _hackward_cos(c: Float[torch.Tensor, "k"], gn: torch.Tensor) -> float:
    """Fraction of the gradient's magnitude that points hack-ward into v_hack:
    ||relu(c)|| / ||g||, where c = V @ g and V rows are orthonormal and oriented
    hack-ward (c_i > 0 means "grad pushes hack-ward on axis i"). In [0, 1].

    relu BEFORE aggregating is the point: the one_sided projection removes only
    relu(c) (the hack-ward axes), and with V orthonormal ||removed|| = ||relu(c)||,
    so this reads directly as "fraction of the grad the projection strips" (a signed
    sum would let +/- axes cancel and read ~0 even while routing a large hack-ward
    magnitude).

    After a one_sided erase, V @ g_proj = min(c, 0) (positive axes zeroed), so
    relu of it is 0 -> cos_post == 0 exactly. That clean SHOULD (cos_post -> 0) is
    the diagnostic; we drop the sign because a one_sided method never acts on the
    safe-ward (negative) part anyway.
    """
    return (torch.relu(c).norm() / gn).item()


@torch.no_grad()
def mean_cos_pre_from_grads(
    grad_dict: dict[str, Float[torch.Tensor, "r"]],
    v_hack: dict[str, Float[torch.Tensor, "k r"]],
) -> float:
    """Mean over modules of ||relu(V @ g)|| / ||g|| (hack-ward fraction, in [0,1]).

    Used to compute per-source cos_pre (cos_pre_s for student-only grad,
    cos_pre_t for teacher-only grad) without mutating model.grad or calling
    the full projection pipeline.
    """
    cs = []
    for name, g in grad_dict.items():
        if g is None or name not in v_hack:
            continue
        V = v_hack[name].to(g.device, dtype=g.dtype)
        gn = g.norm()
        if gn < 1e-12:
            continue
        cs.append(_hackward_cos(V @ g, gn))
    return float(sum(cs) / len(cs)) if cs else float("nan")


def _project_one_module(
    g: Float[torch.Tensor, "r"],
    V: Float[torch.Tensor, "k r"],
    gate_mode: str,
    preserve_magnitude: bool,
    overshoot: float = 1.0,
) -> tuple[Float[torch.Tensor, "r"], Float[torch.Tensor, "r"], float, float, bool]:
    """Per-module top-k removal. Returns (g_proj, removed, cos_pre, cos_post, fired).

    `removed` = overshoot*c_use@V, the vector subtracted from g (computed
    before any preserve_magnitude rescale, so removed ∈ span(V) always).
    Erasure drops it; routing parks it in delta_S_hack. Note g_proj + removed
    == g ONLY when preserve_magnitude is False and overshoot is 1.0; with the
    defaults g_proj is rescaled, so the sum is not the original g (routing does
    not rely on that sum -- see project_delta_S_grad).

    cos_pre / cos_post are the hack-ward FRACTION ||relu(V @ g)|| / ||g|| (in
    [0,1]; see _hackward_cos). cos_pre = how much of the grad points hack-ward
    into v_hack; cos_post = the residual after projection. Under one_sided (and
    no_gate, and overshoot>=1) projection cos_post -> 0 exactly: every hack-ward
    axis was removed, so relu of the residual coefficients is 0.

    `overshoot` scales the removed coefficient: g_proj = g - overshoot*c_use@V.
    overshoot=1.0 just removes the hack-ward component; overshoot=1.1 removes
    110% (a 10% reversal: V@g_proj = -0.1c on fired axes), a milder version of
    gate_mode="reverse" (which is overshoot=2.0 on the full, ungated c).
    """
    gn = g.norm()
    if gn < 1e-12:
        z = torch.zeros_like(g)
        return g, z, 0.0, 0.0, False
    c = V @ g                                       # [k]
    cos_pre = _hackward_cos(c, gn)
    if gate_mode == "no_gate":
        c_use = c
        fired = True
    elif gate_mode == "one_sided":
        mask = (c > 0).to(c.dtype)
        c_use = c * mask
        fired = bool((c_use != 0).any())
    elif gate_mode == "reverse":
        # Subtract 2*c@V: V@g_proj = V@g - 2*(V V^T) c = c - 2c = -c.
        # Flips the sign of the gradient component in span(V); pushes
        # actively away from hack rather than just removing.
        c_use = 2 * c
        fired = True
    else:
        raise ValueError(f"unknown gate_mode={gate_mode!r}")
    if not fired:
        return g, torch.zeros_like(g), cos_pre, cos_pre, False
    removed = overshoot * c_use @ V                 # [r]
    g_proj = g - removed
    gp_n = g_proj.norm()
    if preserve_magnitude and gp_n > 1e-12:
        g_proj = g_proj * (gn / gp_n)
    cos_post = _hackward_cos(V @ g_proj, g_proj.norm().clamp_min(1e-12))
    return g_proj, removed, cos_pre, cos_post, True


@torch.no_grad()
def project_delta_S_grad(
    wrappers: dict,
    v_hack: dict[str, Float[torch.Tensor, "k r"]],
    preserve_magnitude: bool,
    measure_only: bool = False,
    route: bool = False,
    gate_mode: str = "one_sided",
    overshoot: float = 1.0,
) -> dict[str, float]:
    """Per-module top-k removal of hack-aligned grad components.

    For each wrapped module:
      g = delta_S.grad        # [r]
      V = v_hack[name]        # [k, r], rows orthonormal, oriented hack-ward
      c = V @ g               # [k]   per-direction coefficients

    gate_mode="one_sided" (default):
      mask = (c > 0)          # only zap when grad is going hack-ward on that axis
      g' = g - (c * mask) @ V # subtract only positive-coefficient components

    gate_mode="no_gate":
      g' = g - c @ V          # full V·V^T removal, sign-agnostic;
                              # drives ||V g'|| -> 0 exactly. No trust in v_hack
                              # orientation: any motion in span(V) is suspect.

    `preserve_magnitude`: rescale g' to ||g|| after projection.
    `measure_only`: same math, but g is not mutated (the `none` intervention).
    `route`: erase AND park the removed hack-ward component in the quarantine
      knob delta_S_hack.grad (Gradient Routing, Cloud 2410.04332). delta_S gets
      the IDENTICAL g_proj as erase (same gate/preserve/overshoot), so the
      deployment model -- delta_S with delta_S_hack zeroed at eval -- evolves
      under the same update rule as the erase arm (each is its own AdamW param;
      the quarantine's separate optimizer state cannot perturb delta_S). That is
      the sense in which route ⊇ erase: erase == route with the quarantine
      discarded. CAVEAT (not an identity): the combined TRAINING forward
      delta_S + delta_S_hack does NOT reproduce a vanilla update -- AdamW steps
      the two knobs independently, so the sum over-moves hack-ward. That is
      intended (the model keeps hacking during training so the capability lands
      in the quarantine), and it only affects the training trajectory, never the
      ablated deployment. Mutually exclusive with measure_only.

    Diagnostics returned (per call, averaged over modules):
      mean_cos_pre   = mean over modules of ||relu(V @ g)||/||g|| (hack-ward fraction, [0,1])
      mean_cos_post  = same after projection (-> 0 when hack-ward axes were removed)
      frac_fired     = fraction of modules where at least one direction fired (c_i > 0)
    """
    if route and measure_only:
        raise ValueError("route and measure_only are mutually exclusive")
    cos_pre_list, cos_post_list, n_fired = [], [], 0
    for name, info in wrappers.items():
        g = info["delta_S"].grad
        if g is None:
            continue
        if name not in v_hack:  # module dropped by global noise-floor filter
            continue
        V = v_hack[name].to(g.device, dtype=g.dtype)   # [k, r]
        g_proj, removed, cos_pre, cos_post, fired = _project_one_module(
            g, V, gate_mode, preserve_magnitude, overshoot)
        cos_pre_list.append(cos_pre)
        cos_post_list.append(cos_post)
        if fired and not measure_only:
            info["delta_S"].grad = g_proj          # same update rule as erase
            if route:
                # quarantine the discarded hack-ward part; removed ∈ span(V),
                # ablated at eval so its magnitude/overshoot scaling is harmless.
                info["delta_S_hack"].grad = removed
        if fired:
            n_fired += 1
    pre_t = torch.tensor(cos_pre_list); post_t = torch.tensor(cos_post_list)
    return {
        "mean_cos_pre":  pre_t.mean().item(),
        "min_cos_pre":   pre_t.min().item() if pre_t.numel() else float("nan"),
        "max_cos_pre":   pre_t.max().item() if pre_t.numel() else float("nan"),
        "mean_cos_post": post_t.mean().item(),
        "min_cos_post":  post_t.min().item() if post_t.numel() else float("nan"),
        "max_cos_post":  post_t.max().item() if post_t.numel() else float("nan"),
        "frac_fired":    n_fired / len(cos_pre_list) if cos_pre_list else 0.0,
    }