"""
Measurement probes. These use autograd to compute the TRUE gradient purely for
diagnostics -- the learning rules in core.py never call these. The central
quantity is the alignment cos(r_l, -grad_{h_l}L): SDIL claims the innovation
r_l points along the causal descent direction without weight transport.
"""

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

from . import core


def true_hidden_grads(net, x, y):
    """Return (list of g_l = dL/dh_l over hidden layers, scalar loss).
    Uses autograd; temporarily flags W to build the graph, then restores."""
    with torch.enable_grad():
        for W in net.W:
            W.requires_grad_(True)
        x = x.detach()
        h = x
        hs = []
        for l in range(net.L):
            pre = h @ net.W[l].t() + net.b[l]
            if l < net.L - 1:
                act = net.act(pre)
                h = (h + net.res_alpha * act) if (getattr(net, "residual", False) and l >= 1) else act
                hs.append(h)
            else:
                logits = pre
        loss = F.cross_entropy(logits, y)
        grads = torch.autograd.grad(loss, hs, retain_graph=False)
        for W in net.W:
            W.requires_grad_(False)
    return [g.detach() for g in grads], loss.item()


def _row_cos(u, v, eps=1e-12):
    """Mean per-sample cosine between rows of u and v."""
    num = (u * v).sum(1)
    den = u.norm(dim=1) * v.norm(dim=1) + eps
    return (num / den).mean().item()


@torch.no_grad()
def alignment_report(net, x, y, y_onehot, cfg):
    """Per-hidden-layer alignments of the teaching signals with -grad.
    Returns dict of lists (index = hidden layer)."""
    grads, loss = true_hidden_grads(net, x, y)          # g_l
    fwd = net.forward(x)
    h = fwd["h"]
    u = fwd["u"]
    logits = h[-1]
    e = net.output_error(logits, y_onehot)
    c = e

    out = {"cos_r_negg": [], "cos_apical_negg": [], "cos_Ac_negg": [],
           "r_norm": [], "g_norm": [], "loss": loss}
    for l in range(net.L - 1):
        g = grads[l]
        negg = -g
        a = net.apical(l, c, h[l + 1])                  # raw apical (with nuisance)
        Ac = c @ net.A[l].t()                           # pure vectorizer output
        r = a - net.baseline(l, h[l + 1])               # innovation
        # include the phi' gain that the plasticity rule actually applies,
        # so the reported alignment is what the weight update "sees"
        gain = net.act_prime(u[l])
        out["cos_r_negg"].append(_row_cos(r * gain, negg * gain))
        out["cos_apical_negg"].append(_row_cos(a * gain, negg * gain))
        out["cos_Ac_negg"].append(_row_cos(Ac * gain, negg * gain))
        out["r_norm"].append(r.norm(dim=1).mean().item())
        out["g_norm"].append(g.norm(dim=1).mean().item())
    return out


@torch.no_grad()
def loss_decrease_ratio(net, x, y, y_onehot, cfg, step):
    """Single-step descent quality: apply one SDIL update to a scratch copy and
    measure the actual loss drop on the same batch, compared to one plain-SGD
    (true-gradient) step at the same learning rate. Ratio in (0,1] means SDIL
    descends but less steeply than exact GD; <=0 means it went uphill."""
    y_ = y
    L0 = core.loss_ce(net.logits(x), y_).mean().item()

    # SDIL step on a copy
    net_sdil = _clone(net)
    core.sdil_step(net_sdil, x, y_, y_onehot, cfg, step)
    L_sdil = core.loss_ce(net_sdil.logits(x), y_).mean().item()

    # exact-gradient SGD step on a copy at the same eta
    net_bp = _clone(net)
    _sgd_step(net_bp, x, y_, cfg.eta)
    L_bp = core.loss_ce(net_bp.logits(x), y_).mean().item()

    d_sdil = L_sdil - L0
    d_bp = L_bp - L0
    ratio = d_sdil / d_bp if abs(d_bp) > 1e-9 else float("nan")
    return {"dL_sdil": d_sdil, "dL_bp": d_bp, "ratio": ratio}


def _clone(net):
    n = copy.copy(net)
    n.W = [w.clone() for w in net.W]
    n.b = [b.clone() for b in net.b]
    n.A = [a.clone() for a in net.A]
    n.P = [p.clone() for p in net.P]
    n.P_bias = ([p.clone() for p in net.P_bias]
                if getattr(net, "P_bias", None) is not None else None)
    n.Bnuis = [b.clone() for b in net.Bnuis]
    n.mW = [torch.zeros_like(w) for w in net.W]
    n.mb = [torch.zeros_like(b) for b in net.b]
    return n


def _sgd_step(net, x, y, eta):
    with torch.enable_grad():
        for W in net.W:
            W.requires_grad_(True)
        for b in net.b:
            b.requires_grad_(True)
        logits = net.logits(x)
        loss = F.cross_entropy(logits, y)
        grads = torch.autograd.grad(loss, net.W + net.b)
    nW = len(net.W)
    with torch.no_grad():
        for i in range(nW):
            net.W[i] -= eta * grads[i]
            net.b[i] -= eta * grads[nW + i]
    for W in net.W:
        W.requires_grad_(False)
    for b in net.b:
        b.requires_grad_(False)