"""ZBPBlock: an nn.Module wrapper whose backward pass never differentiates through
the wrapped block for the *activation* error.  The returned input gradient is
built purely from (i) function queries of the block (or, in 'oracle' mode, a random
projection of the exact VJP), plus (ii) the exact VJP of an optional *known* skip
path (residual control variate).  Parameter gradients are constructed locally
from the incoming activation error v (J_theta^T v via in-block autograd).
"""
import contextlib
import math
import torch
import torch.nn as nn

from .config import ZBPConfig
from .estimators import (oracle_vjp, zo_vjp, zo_vjp_cd_parts, make_h, control_vjp, Local,
                         causal_score_coordinate_k_vjp, causal_score_coordinate_qk_vjp)

_GENERATORS = {}


def get_generator(device, seed=None):
    key = str(device)
    if key not in _GENERATORS or seed is not None:
        g = torch.Generator(device=device)
        g.manual_seed(1234 if seed is None else seed)
        _GENERATORS[key] = g
    return _GENERATORS[key]


def seed_probes(seed, device="cpu"):
    get_generator(torch.device(device), seed)


class DFA:
    """Holds the output error e = dL/dlogits of the current backward pass for direct feedback alignment.
    Training scripts call DFA.attach(logits) after the forward; blocks in 'dfa' mode read DFA.e.
    DRTP (Frenkel et al. 2021) uses the negated one-hot target instead: set DFA.t via DFA.attach_target(y, n)."""
    e = None
    t = None

    @staticmethod
    def attach_target(y, n_out):
        DFA.t = -torch.nn.functional.one_hot(y, n_out).float()

    @staticmethod
    def attach(logits):
        if logits.requires_grad:
            logits.register_hook(DFA._set)

    @staticmethod
    def _set(g):
        DFA.e = g.detach()


def _feedback_matrix(block, key, shape_in, shape_out, device, dtype):
    """Fixed random feedback map R: shape_out -> shape_in (per feature dims), created once per block."""
    mats = block.__dict__.setdefault("_fb", {})
    if key not in mats:
        g = torch.Generator(device="cpu")
        g.manual_seed(hash((block.name, key)) % (2 ** 31))
        d_in, d_out = math.prod(shape_in), math.prod(shape_out)
        R = torch.randn(d_in, d_out, generator=g) / math.sqrt(d_out)
        mats[key] = R.to(device=device, dtype=dtype)
    return mats[key]


class Recorder:
    """Collects per-block backward diagnostics (exact vs estimated activation error)."""
    _active = None

    def __init__(self):
        self.rows = []

    def __enter__(self):
        Recorder._active = self
        return self

    def __exit__(self, *a):
        Recorder._active = None

    @staticmethod
    def active():
        return Recorder._active

    def add(self, **row):
        self.rows.append(row)

    def summary(self):
        import collections
        by = collections.defaultdict(list)
        for r in self.rows:
            by[r["name"]].append(r)
        out = {}
        for k, rows in by.items():
            out[k] = {
                "cos_mean": sum(r["cos"] for r in rows) / len(rows),
                "relerr_mean": sum(r["relerr"] for r in rows) / len(rows),
                "gnorm": sum(r["gnorm"] for r in rows) / len(rows),
                "ghat_norm": sum(r["ghat_norm"] for r in rows) / len(rows),
                "n": len(rows),
            }
        return out


class ZBPFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, block, *params):
        # Explicit no_grad: no graph through the black box is ever built in the forward pass.
        with torch.no_grad():
            y = block.f(x)
            if block.skip is not None:
                y = y + (x if block.skip_is_identity else block.skip(x))
        ctx.save_for_backward(x)
        ctx.block = block
        return y

    @staticmethod
    def backward(ctx, v):
        (x,) = ctx.saved_tensors
        block = ctx.block
        cfg = block.cfg
        f = block.f
        nbd = block.batch_dims
        v = v.contiguous()
        if block.v_channel is not None:                      # hardware: the held error droops / picks up noise
            v = block.v_channel(v)
        f_params = block._f_params()
        s_params = block._skip_params()
        need_gx = ctx.needs_input_grad[0]
        rec = Recorder.active()
        want_exact = need_gx and (cfg.mode in ("exact", "zero", "noise") or rec is not None
                                  or block.probe_slice is not None or block.score_probe is not None)
        if cfg.mode == "replay":
            want_exact = False
        if cfg.mode in ("dfa", "opzo", "drtp") and block.skip is not None and not block.skip_is_identity and need_gx:
            pass  # DFA ignores the skip path entirely (standard DFA); skip params still get their local gradient

        # --- local parameter gradients (and exact g_x if needed) via in-block autograd ---
        g_exact = None
        g_params = [None] * len(f_params)
        if want_exact or (f_params and cfg.param_mode == "local"):
            with torch.enable_grad():
                xd = x.detach().requires_grad_(want_exact)
                y = f(xd)
                inputs = ([xd] if want_exact else []) + list(f_params)
                if inputs:
                    grads = torch.autograd.grad(y, inputs, v, allow_unused=True)
                    if want_exact:
                        g_exact, grads = grads[0], grads[1:]
                    g_params = list(grads)

        # --- activation error through the black-box branch ---
        g_hat = None
        queries = 0.0
        if need_gx:
            gen = get_generator(x.device)
            if cfg.mode == "exact":
                g_hat = g_exact
            elif cfg.mode == "replay":
                g_hat = None                                  # full input error injected below
            elif cfg.mode in ("zero", "noise"):
                g_hat, queries = control_vjp(g_exact, cfg, gen, nbd)
            elif cfg.mode == "fa":
                # feedback alignment: J_f^T v replaced by a fixed random linear map of the incoming error.
                # Sequence models ([B, T, d]): one d_in x d_out map shared across tokens (standard practice);
                # otherwise a full map over the flattened feature dims.
                if x.dim() == 3 and v.dim() == 3 and x.shape[1] == v.shape[1]:
                    R = _feedback_matrix(block, "fa_tok", x.shape[2:], v.shape[2:], x.device, x.dtype)
                    g_hat = (v @ R.t()) * block.fb_scale
                else:
                    feat_in, feat_out = x.shape[nbd:], v.shape[nbd:]
                    R = _feedback_matrix(block, "fa", feat_in, feat_out, x.device, x.dtype)
                    g_hat = (v.reshape(*v.shape[:nbd], -1) @ R.t()).reshape(x.shape) * block.fb_scale
            elif cfg.mode == "opzo":
                # OPZO-style baseline (Xiao et al. 2024): DFA topology with a *learned* feedback matrix — the EMA of a
                # zeroth-order estimate of the data-averaged Jacobian^T from this block's input to the logits,
                # obtained from perturbed forward passes of the whole network (see opzo_update).  Biased by design.
                e = DFA.e
                if e is None:
                    raise RuntimeError("DFA.e is unset: call DFA.attach(logits) after the forward pass")
                M = block.__dict__.get("_opzo", {}).get("M")
                if M is None:
                    g_hat = torch.zeros_like(x)
                elif e.dim() == 3 and x.dim() >= 3 and e.shape[1] == x.shape[1]:      # per-token feedback
                    g_hat = (e @ M.t()).reshape(x.shape)
                else:                                                                  # per-sample feedback
                    g_hat = (e.reshape(e.shape[0], -1) @ M.t()).reshape(x.shape)
            elif cfg.mode in ("dfa", "drtp"):
                # direct feedback alignment: the output error is projected straight to this block's input;
                # DRTP: the (negated) one-hot target takes the place of the error
                e = DFA.e if cfg.mode == "dfa" else DFA.t
                if e is None:
                    raise RuntimeError("DFA.e / DFA.t is unset: call DFA.attach(logits) / DFA.attach_target(y, n) first")
                e = e.to(x.dtype)
                if e.dim() == 3 and x.dim() >= 3 and e.shape[1] == x.shape[1]:  # per-token error [B, T, V]
                    R = _feedback_matrix(block, "dfa_tok", x.shape[2:], e.shape[2:], x.device, x.dtype)
                    g_hat = (e @ R.t()).reshape(x.shape) * block.fb_scale
                else:                                                            # one error vector per sample
                    R = _feedback_matrix(block, "dfa", x.shape[1:], e.shape[1:], x.device, x.dtype)
                    g_hat = (e.reshape(e.shape[0], -1) @ R.t()).reshape(x.shape) * block.fb_scale
            elif block.score_probe is not None:
                # Probe the causal softmax score rows and map their VJP to K through QK^T.
                # Other input slices retain their exact diagnostic VJPs (notably A^T delta for V).
                spec = block.score_probe
                g_hat = g_exact.clone()
                a = spec["d"]
                if spec.get("qk"):     # Q and K both from the score-VJP; V exact via the known attention weights
                    gq, gk, queries = causal_score_coordinate_qk_vjp(
                        x, v, cfg, gen, spec["d"], spec["heads"], spec["dk"], spec["dv"], spec["window"])
                    g_hat[..., :a] = gq
                else:
                    gk, queries = causal_score_coordinate_k_vjp(
                        x, v, cfg, gen, spec["d"], spec["heads"], spec["dk"], spec["dv"], spec["window"])
                g_hat[..., a:a + spec["dk"]] = gk
            elif block.probe_slice is not None:
                # ZO estimate only for coordinates [a, b) of the last dim (with its own unit structure);
                # the other coordinates keep the exact Jacobian (diagnostic: which input is hard to probe)
                a, b = block.probe_slice
                sbd = block.slice_batch_dims or nbd
                xs = x[..., a:b].contiguous()
                def f_slice(xsl):
                    k = xsl.shape[0] // x.shape[0]
                    xx = x.repeat(k, *([1] * (x.dim() - 1))) if k > 1 else x
                    return f(torch.cat([xx[..., :a], xsl, xx[..., b:]], dim=-1))
                if cfg.mode == "oracle":
                    gs, queries = oracle_vjp(f_slice, xs, v, cfg, gen, sbd, None)
                else:
                    h = make_h(f_slice, v, sbd, None, cfg.readout_noise, gen)
                    gs, queries = zo_vjp(h, xs, cfg, gen, sbd, None)
                g_hat = g_exact.clone()
                g_hat[..., a:b] = gs
            elif cfg.mode == "oracle":
                g_hat, queries = oracle_vjp(f, x, v, cfg, gen, nbd, block.local)
            elif cfg.mode == "cd" and cfg.surrogate and block.local is None:
                # learned linear control variate: g_hat = g_raw - s * (P_u g_sur - g_sur), g_sur = A v,
                # with A and s fitted on *previous* steps' data only (unbiased conditional on the past)
                h = make_h(f, v, nbd, None, cfg.readout_noise, gen)
                cyc = block.stats["backward_calls"] if cfg.probe == "hadamard_cycle" else None
                u, D, g_raw, queries = zo_vjp_cd_parts(h, x, cfg, gen, nbd, cyc)
                g_hat = block._apply_surrogate(x, v, u, D, g_raw, nbd)
            else:
                fq = f if block.measure is None else (lambda xp: block.measure(f(xp)))
                h = make_h(fq, v, nbd, block.local, cfg.readout_noise, gen)
                cyc = block.stats["backward_calls"] if cfg.probe == "hadamard_cycle" else None
                g_hat, queries = zo_vjp(h, x, cfg, gen, nbd, block.local, cyc)
            block.stats["queries"] += queries
            block.stats["backward_calls"] += 1
            if rec is not None and g_hat is not None:
                ge = g_exact.reshape(g_exact.shape[0], -1)
                gh = g_hat.reshape(g_hat.shape[0], -1)
                cos = torch.nn.functional.cosine_similarity(ge, gh, dim=1)
                rel = (gh - ge).norm(dim=1) / ge.norm(dim=1).clamp_min(1e-30)
                rec.add(name=block.name, cos=cos.mean().item(), relerr=rel.mean().item(),
                        gnorm=ge.norm(dim=1).mean().item(), ghat_norm=gh.norm(dim=1).mean().item(),
                        vnorm=v.reshape(v.shape[0], -1).norm(dim=1).mean().item(), queries=queries,
                        # batch-level (mean over samples) cosine, what parameter updates "see"
                        cos_batchmean=torch.nn.functional.cosine_similarity(
                            ge.mean(0, keepdim=True), gh.mean(0, keepdim=True), dim=1).item())

        # --- exact skip path (residual control variate) ---
        g_skip_params = [None] * len(s_params)
        grad_x = None
        if need_gx and cfg.mode == "replay":
            grad_x = block.replay
            if s_params:
                with torch.enable_grad():
                    xd = x.detach()
                    s = block.skip(xd)
                    g_skip_params = list(torch.autograd.grad(s, list(s_params), v, allow_unused=True))
        elif need_gx and cfg.mode in ("dfa", "opzo", "drtp"):
            grad_x = g_hat                                    # DFA/DRTP/OPZO: direct projection replaces the whole backward
            if s_params:                                      # skip-path parameters still get their local gradient
                with torch.enable_grad():
                    xd = x.detach()
                    s = block.skip(xd)
                    g_skip_params = list(torch.autograd.grad(s, list(s_params), v, allow_unused=True))
        elif need_gx:
            if block.skip is None:
                grad_x = g_hat
            elif block.skip_is_identity:
                grad_x = v + g_hat
            else:
                with torch.enable_grad():
                    xd = x.detach().requires_grad_(True)
                    s = block.skip(xd)
                    grads = torch.autograd.grad(s, [xd] + list(s_params), v, allow_unused=True)
                grad_x = grads[0] + g_hat
                g_skip_params = list(grads[1:])
        elif s_params:
            with torch.enable_grad():
                xd = x.detach()
                s = block.skip(xd)
                g_skip_params = list(torch.autograd.grad(s, list(s_params), v, allow_unused=True))
        if block.capture and grad_x is not None:
            block.captured = grad_x.detach()
        return (grad_x, None, *g_params, *g_skip_params)


class ZBPBlock(nn.Module):
    """y = f(x)  or  y = skip(x) + f(x).

    The activation error through f is estimated with zeroth-order queries (or the oracle
    projection); the skip path (identity or a cheap differentiable module) is propagated
    exactly.  With cfg.mode == 'bp' the block is ordinary autograd.
    """

    def __init__(self, f, cfg=None, skip=None, batch_dims=1, name="", estimate_input_grad=True, local=None):
        super().__init__()
        self.f = f
        # local=(radius, stride): conv-like block whose output position p depends only on input
        # positions within `radius` of stride*p -> per-position measurements (lower variance)
        self.local = Local(*local) if isinstance(local, (tuple, list)) else local
        self.cfg = cfg if cfg is not None else ZBPConfig()
        self.batch_dims = batch_dims
        self.name = name
        self.estimate_input_grad = estimate_input_grad
        if skip is None:
            self.skip = None
            self.skip_is_identity = False
        elif skip == "identity" or isinstance(skip, nn.Identity):
            self.skip = nn.Identity()
            self.skip_is_identity = True
        else:
            self.skip = skip
            self.skip_is_identity = False
        self.stats = {"queries": 0.0, "backward_calls": 0}
        self.fb_scale = 1.0       # scale of the fixed random feedback map in 'fa' / 'dfa' modes
        self.probe_slice = None   # (a, b): ZO-probe only last-dim coordinates a:b, exact elsewhere
        self.slice_batch_dims = None
        self.score_probe = None   # metadata for causal score-coordinate K probing
        self.capture = False      # store the full input error of the next backward in .captured
        self.captured = None
        self.replay = None        # 'replay' mode returns this tensor as the input error
        self.measure = None       # hardware measurement channel applied to y inside the query readout
        self.v_channel = None     # hardware error-transport channel (sample/hold) applied to the incoming v

    def _f_params(self):
        return [p for p in self.f.parameters() if p.requires_grad]

    # ---- learned linear control variate (prior-free): A_t = EMA[g_raw v^T] EMA[v v^T]^-1, s_t from EMA[D P]/EMA[P^2] ----
    def _apply_surrogate(self, x, v, u, D, g_raw, nbd):
        cfg = self.cfg
        n = u.shape[0]
        bs = x.shape[:nbd]
        B = math.prod(bs) if len(bs) else 1
        xf, vf, gf = x.reshape(B, -1), v.reshape(B, -1), g_raw.reshape(B, -1)
        uf, Df = u.reshape(n, B, -1), D.reshape(n, B)
        din, dout = xf.shape[1], vf.shape[1]
        st = self.__dict__.setdefault("_sur", None)
        if st is None:
            st = {"S_gv": torch.zeros(din, dout, device=x.device, dtype=x.dtype),
                  "S_vv": torch.zeros(dout, dout, device=x.device, dtype=x.dtype),
                  "A": None, "s": 0.0, "ema_dp": 0.0, "ema_pp": 0.0, "cnt": 0}
            self.__dict__["_sur"] = st
        ema = cfg.surrogate_ema
        g_hat = gf
        if st["A"] is not None:
            with torch.no_grad():
                g_sur = vf @ st["A"].t()                              # [B, din]  surrogate  A v
                P = (uf * g_sur[None]).sum(-1)                        # [n, B]    u_i^T g_sur (digital)
                g_sur_raw = (uf * P[..., None]).mean(0)               # same-probe projection of the surrogate
                if st["s"] > 0:
                    g_hat = gf - st["s"] * (g_sur_raw - g_sur)        # E[g_hat] = g for any A, s fixed before this step
                dp, pp = (Df * P).sum().item(), (P * P).sum().item()
                st["ema_dp"] = ema * st["ema_dp"] + (1 - ema) * dp
                st["ema_pp"] = ema * st["ema_pp"] + (1 - ema) * pp
                st["s"] = float(min(1.0, max(0.0, st["ema_dp"] / max(st["ema_pp"], 1e-30))))
        with torch.no_grad():                                          # regression statistics for the NEXT steps
            st["S_gv"] = ema * st["S_gv"] + (1 - ema) * (gf.t() @ vf) / B
            st["S_vv"] = ema * st["S_vv"] + (1 - ema) * (vf.t() @ vf) / B
            st["cnt"] += 1
            if st["cnt"] >= 2:
                ridge = cfg.surrogate_ridge * st["S_vv"].diagonal().mean().clamp_min(1e-30)
                M = st["S_vv"] + ridge * torch.eye(dout, device=x.device, dtype=x.dtype)
                st["A"] = torch.linalg.solve(M, st["S_gv"].t()).t()   # [din, dout]: least squares of g_raw on v
        self.stats["surrogate_s"] = st["s"]
        return g_hat.reshape(g_raw.shape)

    def _skip_params(self):
        if self.skip is None or self.skip_is_identity:
            return []
        return [p for p in self.skip.parameters() if p.requires_grad]

    def forward(self, x):
        pert = self.__dict__.get("opzo_perturb")
        if pert is not None:                                   # OPZO estimation pass: perturb this block's input
            alpha, gen = pert
            z = torch.randn(x.shape, generator=gen, device=x.device, dtype=x.dtype)
            self.__dict__.setdefault("_opzo", {"M": None})["z"] = z
            x = x + alpha * z
        if self.cfg.mode == "bp" or not torch.is_grad_enabled():
            y = self.f(x)
            if self.skip is not None:
                y = y + (x if self.skip_is_identity else self.skip(x))
            return y
        if not self.estimate_input_grad:
            x = x.detach()
        return ZBPFunction.apply(x, self, *self._f_params(), *self._skip_params())

    def extra_repr(self):
        loc = f", local=({self.local.radius},{self.local.stride})" if self.local is not None else ""
        return f"name={self.name}, mode={self.cfg.mode}, n={self.cfg.n_probes}, skip={'id' if self.skip_is_identity else (self.skip is not None)}{loc}"


def zbp_blocks(model):
    return [m for m in model.modules() if isinstance(m, ZBPBlock)]


def opzo_update(model, forward_fn, out_clean, alpha, lam, gen):
    """One OPZO estimation step: perturb the inputs of all 'opzo' blocks at once (alpha * Gaussian), run one extra
    forward of the whole network, and update each block's feedback matrix
        M <- lam * M + (1 - lam) * mean_units z (delta_logits / alpha)^T
    (cross-block responses are zero-mean noise, as in OPZO's single noisy pass).  Per-token when the logits are
    [B, T, V] and the block input has a matching token dim, per-sample otherwise."""
    blocks = [b for b in zbp_blocks(model) if b.cfg.mode == "opzo"]
    if not blocks:
        return
    for b in blocks:
        b.__dict__["opzo_perturb"] = (alpha, gen)
    with torch.no_grad():
        out_p = forward_fn()
    for b in blocks:
        b.__dict__["opzo_perturb"] = None
    delta = (out_p - out_clean.detach()) / alpha
    with torch.no_grad():
        for b in blocks:
            st = b.__dict__["_opzo"]
            z = st.pop("z")
            if delta.dim() == 3 and z.dim() >= 3 and z.shape[1] == delta.shape[1]:
                zf = z.reshape(z.shape[0], z.shape[1], -1)
                M_new = torch.einsum("btd,btv->dv", zf, delta) / (z.shape[0] * z.shape[1])
            else:
                M_new = z.reshape(z.shape[0], -1).t() @ delta.reshape(delta.shape[0], -1) / z.shape[0]
            st["M"] = M_new if st["M"] is None else lam * st["M"] + (1 - lam) * M_new
            b.stats["queries"] += 1.0


def set_mode(model, cfg=None, **kw):
    """Replace / update the estimator config of every ZBPBlock in a model."""
    for b in zbp_blocks(model):
        b.cfg = (cfg if cfg is not None else b.cfg).replace(**kw)


def total_queries(model):
    return sum(b.stats["queries"] for b in zbp_blocks(model))


def reset_queries(model):
    for b in zbp_blocks(model):
        b.stats = {"queries": 0.0, "backward_calls": 0}