"""
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):
    """Mean per-sample cosine, excluding rows with an exact zero vector.

    Adding an absolute epsilon to the denominator makes cosine spuriously
    depend on signal magnitude near convergence.  In particular, a positive
    per-sample rescaling can then appear to change direction.  Exact zero rows
    have no defined direction and are omitted instead.
    """
    num = (u * v).sum(1)
    den = u.norm(dim=1) * v.norm(dim=1)
    valid = den > 0
    if not valid.any():
        return float("nan")
    return (num[valid] / den[valid]).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]
    context = h[-2]
    e = net.output_error(logits, y_onehot)
    c = e

    out = {"cos_r_negg": [], "cos_innovation_negg": [],
           "cos_apical_negg": [], "cos_Ac_negg": [],
           "r_norm": [], "g_norm": [], "traffic_norm": [],
           "traffic_residual_norm": [], "traffic_r2": [], "loss": loss}
    for l in range(net.L - 1):
        g = grads[l]
        negg = -g
        teaching, a, innovation = core.teaching_signal(
            net, l, c, h[l + 1], cfg, context)
        Ac = net.vectorizer(l, c, h[l + 1], context)    # pure vectorizer output
        traffic = net.apical_traffic(l, h[l + 1], context)
        unexplained = traffic - net.baseline(l, h[l + 1])
        # 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(teaching * gain, negg * gain))
        out["cos_innovation_negg"].append(_row_cos(innovation * 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(teaching.norm(dim=1).mean().item())
        out["g_norm"].append(g.norm(dim=1).mean().item())
        out["traffic_norm"].append(traffic.norm(dim=1).mean().item())
        out["traffic_residual_norm"].append(unexplained.norm(dim=1).mean().item())
        traffic_power = traffic.pow(2).mean()
        if traffic_power > 0:
            out["traffic_r2"].append(
                (1.0 - unexplained.pow(2).mean() / traffic_power).item())
        else:
            out["traffic_r2"].append(float("nan"))
    return out


@torch.no_grad()
def fa_alignment_report(net, x, y, y_onehot):
    """Alignment of sequential-FA local update signals with exact descent.

    ``FANet.fa_hidden_gradients`` performs the same fixed-feedback transport
    used during learning.  Autograd appears only here to obtain the reference
    hidden-state gradients for measurement.
    """
    grads, loss = true_hidden_grads(net, x, y)
    fwd = net.forward(x)
    estimated = net.fa_hidden_gradients(fwd, y_onehot)
    cosines = []
    for l, (estimate, true_grad) in enumerate(zip(estimated, grads)):
        gain = net.act_prime(fwd["u"][l])
        cosines.append(_row_cos(-estimate * gain, -true_grad * gain))
    return {"cos_fa_negg": cosines, "loss": loss}


@torch.no_grad()
def nodepert_alignment_report(net, x, y, cfg):
    """Alignment of the unamortized perturbation signal actually used."""
    grads, loss = true_hidden_grads(net, x, y)
    estimator = (core.simultaneous_node_perturbation_targets
                 if cfg.pert_mode == "simultaneous"
                 else core.node_perturbation_targets)
    targets = estimator(net, x, y, sigma=cfg.pert_sigma, n_dirs=cfg.pert_ndirs)
    fwd = net.forward(x)
    cosines = []
    q_norm = []
    g_norm = []
    for layer, (target, grad) in enumerate(zip(targets, grads)):
        gain = net.act_prime(fwd["u"][layer])
        cosines.append(_row_cos(target * gain, -grad * gain))
        q_norm.append(target.norm(dim=1).mean().item())
        g_norm.append(grad.norm(dim=1).mean().item())
    return {"cos_q_negg": cosines, "q_norm": q_norm, "g_norm": g_norm, "loss": loss}


@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.A_gate = ([a.clone() for a in net.A_gate]
                if getattr(net, "A_gate", None) is not None else None)
    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.Bcontext = [b.clone() for b in net.Bcontext]
    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)