"""Zeroth-order estimators of g = J^T v for a black-box block.

Conventions
-----------
A block maps x (shape [*batch_shape, *feat]) to y.  A *measurement* function
`h(x_pert)` returns v^T f(x_pert) per measurement unit:
  * global units:  h -> [*batch_shape]              (sum over all feature dims)
  * local units:   h -> [*batch_shape, H', W']      (sum over channels only; conv blocks)
`h` accepts any integer multiple of the batch stacked along dim 0, so several
probes can be evaluated in one call.

For local units the directional derivatives D_p (one per output position p) are
box-summed over the receptive field before multiplying the probe, which keeps the
estimator unbiased while removing the variance contributed by far-away positions.

Every estimator returns (g_hat, queries) where queries is the number of block
evaluations *per sample unit* that a physical system would have to perform.
"""
import math
import torch
import torch.nn.functional as F

from .probes import sample_probes


class Local:
    """Receptive-field geometry of a conv-like block: output position p depends on input
    positions q with |q - stride*p|_inf <= radius (input coordinates)."""

    def __init__(self, radius, stride=1):
        self.radius, self.stride = int(radius), int(stride)

    def gather(self, D, in_hw):
        """D: [n, *bs, H', W'] per-output-position derivatives -> S: [n, *bs, H, W] box sums."""
        lead = D.shape[:-2]
        Hp, Wp = D.shape[-2:]
        H, W = in_hw
        k = 2 * self.radius + 1
        Dm = D.reshape(-1, 1, Hp, Wp)
        ones = torch.ones(1, 1, k, k, dtype=D.dtype, device=D.device)
        oph = H - ((Hp - 1) * self.stride + 1)
        opw = W - ((Wp - 1) * self.stride + 1)
        S = F.conv_transpose2d(Dm, ones, stride=self.stride, padding=self.radius, output_padding=(oph, opw))
        return S.reshape(*lead, H, W)


class Local1D:
    """Causal token-window locality for [B, T, d] blocks: output token t depends on input tokens
    s in [t-window+1, t].  Measurement units = output tokens; gather = causal window box-sum."""
    kind = "1d"

    def __init__(self, window):
        self.window = int(window)

    def gather(self, D, in_shape):
        """D: [n, *bs, T] per-output-token derivatives -> S: [n, *bs, T], S[s] = sum_{t=s}^{s+W-1} D[t]."""
        T = D.shape[-1]
        c = torch.cumsum(D, dim=-1)
        pad = torch.zeros_like(c[..., :1])
        c0 = torch.cat([pad, c], dim=-1)                                # c0[t] = sum_{<t}
        hi = torch.clamp(torch.arange(T, device=D.device) + self.window, max=T)
        lo = torch.arange(T, device=D.device)
        return c0[..., hi] - c0[..., lo]


def _probes_for(x, n_probes, family, gen, n_batch_dims, cycle=None):
    bs = x.shape[:n_batch_dims]
    feat = x.shape[n_batch_dims:]
    d = math.prod(feat)
    u = sample_probes(n_probes, tuple(bs), d, family, gen, x.device, x.dtype, cycle=cycle)
    return u.reshape(n_probes, *bs, *feat), d


def zo_vjp_cd_parts(h, x, cfg, gen, n_batch_dims=1, cycle=None):
    """Central-difference measurements together with their probes (global units only):
    returns (u [n,*bs,*feat], D [n,*bs] = u_i^T g + O(eps^2), g_hat, queries)."""
    u, d = _probes_for(x, cfg.n_probes, cfg.probe, gen, n_batch_dims, cycle)
    eps = _eps_for(cfg, d)
    D = _eval_pairs(h, x, u, [eps], n_batch_dims, cfg.batch_probes, cfg.probe_chunk)[0]
    return u, D, _combine(u, D, n_batch_dims, None, x.shape), float(2 * cfg.n_probes)


def _combine(u, D, n_batch_dims, local, in_shape):
    """g_hat = (1/n) sum_i u_i * S_i, where S_i are the (gathered) directional derivatives."""
    n = u.shape[0]
    if local is None:
        S = D.reshape(*D.shape, *([1] * (u.dim() - 1 - n_batch_dims)))
    elif getattr(local, "kind", "2d") == "1d":
        S = local.gather(D, in_shape).unsqueeze(-1)           # [n, *bs, T, 1] broadcast over features
    else:
        S = local.gather(D, in_shape[-2:])                   # [n, *bs, H, W]
        S = S.unsqueeze(n_batch_dims + 1)                     # broadcast over channels: [n, *bs, 1, H, W]
    return (u * S).mean(0)


def _eps_for(cfg, d):
    return cfg.eps if cfg.eps_mode == "coord" else cfg.eps / math.sqrt(d)


def _eval_pairs(h, x, u, eps_list, n_batch_dims, batch_probes=True, chunk=0):
    """D(eps, u_i) = (h(x+eps u_i) - h(x-eps u_i)) / (2 eps) for all probes/eps -> [len(eps), n, *bs, *meas]."""
    n = u.shape[0]
    B0 = x.shape[0]
    outs = []
    if batch_probes:
        step = n if chunk <= 0 else max(1, min(chunk, n))
        for eps in eps_list:
            parts = []
            for i0 in range(0, n, step):
                xs = []
                for i in range(i0, min(n, i0 + step)):
                    xs.append(x + eps * u[i])
                    xs.append(x - eps * u[i])
                hv = h(torch.cat(xs, 0))
                hv = hv.reshape(len(xs) // 2, 2, B0, *hv.shape[1:])
                parts.append((hv[:, 0] - hv[:, 1]) / (2 * eps))
            outs.append(torch.cat(parts, 0))
    else:
        for eps in eps_list:
            outs.append(torch.stack([(h(x + eps * u[i]) - h(x - eps * u[i])) / (2 * eps) for i in range(n)]))
    return torch.stack(outs)


def exact_directional(f, x, v, u, n_batch_dims=1, local=None, chunk=16):
    """Oracle: exact directional derivatives D_i = <v, J u_i> per measurement unit via forward-mode AD
    (vmapped over probes in chunks; falls back to a loop if the block is not vmap-compatible)."""
    def one(t):
        return torch.func.jvp(f, (x,), (t,))[1]
    outs = []
    n = u.shape[0]
    try:
        for i in range(0, n, chunk):
            outs.append(torch.func.vmap(one)(u[i:i + chunk]))
        Ju = torch.cat(outs, 0)
    except Exception:
        Ju = torch.stack([one(u[i]) for i in range(n)])
    prod = Ju * v.unsqueeze(0)
    if local is None:
        return prod.reshape(n, *x.shape[:n_batch_dims], -1).sum(-1)
    if getattr(local, "kind", "2d") == "1d":
        return prod.sum(-1)
    return prod.sum(n_batch_dims + 1)                         # sum over channel dim only


def oracle_vjp(f, x, v, cfg, gen, n_batch_dims=1, local=None):
    """Stage A: (1/n) sum_i u_i <v, J u_i>  (== (1/n) sum_i u_i (u_i^T g)).  Conditionally unbiased."""
    u, d = _probes_for(x, cfg.n_probes, cfg.probe, gen, n_batch_dims)
    with torch.no_grad():
        D = exact_directional(f, x, v, u, n_batch_dims, local)
    return _combine(u, D, n_batch_dims, local, x.shape), 0.0


def oracle_projection(g, cfg, gen, n_batch_dims=1):
    """Projection form of the oracle for a given exact VJP g (global units only)."""
    u, d = _probes_for(g, cfg.n_probes, cfg.probe, gen, n_batch_dims)
    n = u.shape[0]
    bs = g.shape[:n_batch_dims]
    D = (u.reshape(n, *bs, -1) * g.reshape(1, *bs, -1)).sum(-1)
    return _combine(u, D, n_batch_dims, None, g.shape), 0.0


def zo_vjp(h, x, cfg, gen, n_batch_dims=1, local=None, cycle=None):
    """Function-query estimate of J^T v.  Returns (g_hat, queries_per_sample_unit)."""
    if cfg.mode == "cd" and cfg.probe == "coordinate_all" and local is None:
        # analytic combine: u_i = sqrt(d) e_i  ->  g_hat[..., i] = D_i / sqrt(d); the probe tensor is a broadcast
        # view of eye(d), so memory stays O(d^2 + measurements) instead of O(n * batch * d)
        bs = x.shape[:n_batch_dims]
        d = math.prod(x.shape[n_batch_dims:])
        if cfg.n_probes != d:
            raise ValueError(f"coordinate_all needs n == d (n={cfg.n_probes}, d={d})")
        eps = _eps_for(cfg, d)
        ident = math.sqrt(d) * torch.eye(d, device=x.device, dtype=x.dtype)
        uv = ident.reshape(d, *([1] * len(bs)), d).expand(d, *bs, d)
        D = _eval_pairs(h, x, uv, [eps], n_batch_dims, cfg.batch_probes, cfg.probe_chunk)[0]
        return (D.movedim(0, -1) / math.sqrt(d)).reshape(x.shape), float(2 * d)
    if cfg.mode == "cd" and cfg.probe_chunk and cfg.n_probes > 4 * cfg.probe_chunk and cfg.probe != "hadamard_cycle":
        # probe-chunked sample -> evaluate -> accumulate: memory O(chunk * batch * d) for any n
        bs = x.shape[:n_batch_dims]
        d = math.prod(x.shape[n_batch_dims:])
        eps = _eps_for(cfg, d)
        g_acc = torch.zeros_like(x)
        done = 0
        while done < cfg.n_probes:
            c = min(cfg.probe_chunk, cfg.n_probes - done)
            uc, _ = _probes_for(x, c, cfg.probe, gen, n_batch_dims)
            D = _eval_pairs(h, x, uc, [eps], n_batch_dims, cfg.batch_probes, cfg.probe_chunk)[0]
            g_acc += _combine(uc, D, n_batch_dims, local, x.shape) * c
            done += c
        return g_acc / cfg.n_probes, float(2 * cfg.n_probes)
    u, d = _probes_for(x, cfg.n_probes, cfg.probe, gen, n_batch_dims, cycle)
    n = cfg.n_probes
    eps = _eps_for(cfg, d)
    bs = x.shape[:n_batch_dims]
    B0 = x.shape[0]

    if cfg.mode == "forward":
        if cfg.batch_probes:
            hv = h(torch.cat([x] + [x + eps * u[i] for i in range(n)], 0))
            hv = hv.reshape(n + 1, B0, *hv.shape[1:])
            D = (hv[1:] - hv[:1]) / eps
        else:
            h0 = h(x)
            D = torch.stack([(h(x + eps * u[i]) - h0) / eps for i in range(n)])
        return _combine(u, D, n_batch_dims, local, x.shape), float(n + 1)

    if cfg.mode == "cd":
        D = _eval_pairs(h, x, u, [eps], n_batch_dims, cfg.batch_probes, cfg.probe_chunk)[0]
        return _combine(u, D, n_batch_dims, local, x.shape), float(2 * n)

    if cfg.mode == "richardson":
        D = _eval_pairs(h, x, u, [eps, eps / 2], n_batch_dims, cfg.batch_probes, cfg.probe_chunk)
        R = (4 * D[1] - D[0]) / 3
        return _combine(u, R, n_batch_dims, local, x.shape), float(4 * n)

    if cfg.mode in ("ml", "ml_richardson"):
        rich = cfg.mode == "ml_richardson"
        if cfg.level_sampling == "per_sample":
            U = torch.rand((n, *bs), generator=gen, device=x.device, dtype=torch.float64)
        else:
            U = torch.rand((n,), generator=gen, device=x.device, dtype=torch.float64)
            U = U.reshape(n, *([1] * len(bs))).expand(n, *bs)
        N = torch.floor(-torch.log2(U) / cfg.alpha).clamp(max=cfg.max_level).to(torch.int64)
        if cfg.eps_min > 0:   # truncate the ladder where readout noise would dominate (bias O(eps_min^2))
            kmax = max(0, int(math.floor(math.log2(eps / cfg.eps_min)))) if eps > cfg.eps_min else 0
            N = N.clamp(max=kmax)
        Nmax = int(N.max().item())
        K = Nmax + (2 if rich else 1)
        eps_list = [eps * 2.0 ** (-k) for k in range(K)]
        D = _eval_pairs(h, x, u, eps_list, n_batch_dims, cfg.batch_probes, cfg.probe_chunk)   # [K, n, *bs, *meas]
        S = (4 * D[1:] - D[:-1]) / 3 if rich else D
        est = S[0].clone()
        extra = S.dim() - 1 - N.dim()
        for k in range(1, Nmax + 1):
            mask = (N >= k).to(x.dtype).reshape(*N.shape, *([1] * extra))
            est = est + mask * (S[k] - S[k - 1]) / (2.0 ** (-cfg.alpha * k))
        per = 2.0 * (N.to(torch.float64) + 1 + (1 if rich else 0))
        queries = float(per.sum().item() / (math.prod(bs) if len(bs) else 1))
        return _combine(u, est, n_batch_dims, local, x.shape), queries

    raise ValueError(f"unknown ZBP mode {cfg.mode!r}")


def causal_score_coordinate_k_vjp(qkv, v_out, cfg, gen, d, n_heads, dk=None, dv=None, window=None):
    """Estimate attention's K-side VJP by probing causal score rows.

    The black-box interface here is ``softmax(scores) @ value`` rather than K itself.
    One paired query perturbs one valid score coordinate in every (sample, head,
    output-token) row, and the output supplies a separate scalar measurement for
    every such row.  For a row with ``m`` valid keys, the estimator visits all
    coordinates when ``m <= n``; otherwise it samples ``n`` coordinates without
    replacement and applies the Horvitz--Thompson weight ``m / n``.  Thus the
    row-gradient relative MSE is exactly ``max(m / n - 1, 0)`` in the directional-
    derivative limit, instead of ``(m - 1) / n`` for i.i.d. dense probes.

    The estimated score VJP is mapped to K through the digitally known bilinear
    score map QK^T/sqrt(dh).  Q and V are held fixed during the function queries.
    Returns ``(g_k, queries_per_sequence)``; full-attention K is [B,T,d], while
    multi-query K is [B,T,dh].
    """
    if cfg.mode not in ("cd", "oracle"):
        raise ValueError(f"causal score-coordinate probing supports cd/oracle, got {cfg.mode!r}")
    B, T, _ = qkv.shape
    H = int(n_heads)
    dh = int(d) // H
    dk = int(d if dk is None else dk)
    dv = int(d if dv is None else dv)
    if int(d) % H:
        raise ValueError(f"attention width {d} is not divisible by {H} heads")
    if dk not in (int(d), dh) or dv not in (int(d), dh):
        raise ValueError("score-coordinate probing supports full or multi-query K/V layouts")

    q0, k0, value0 = qkv.split([int(d), dk, dv], dim=-1)
    q = q0.view(B, T, H, dh).transpose(1, 2)
    if dk == dh:
        k = k0.view(B, T, 1, dh).transpose(1, 2).expand(B, H, T, dh)
    else:
        k = k0.view(B, T, H, dh).transpose(1, 2)
    if dv == dh:
        value = value0.view(B, T, 1, dh).transpose(1, 2).expand(B, H, T, dh)
    else:
        value = value0.view(B, T, H, dh).transpose(1, 2)
    delta = v_out.view(B, T, H, dh).transpose(1, 2)

    idx = torch.arange(T, device=qkv.device)
    valid = idx[None, :] <= idx[:, None]
    if window is not None:
        valid = valid & (idx[None, :] > idx[:, None] - int(window))
    counts = valid.sum(-1)                                      # [T], valid coordinates per row
    scores = (q @ k.transpose(-2, -1)) / math.sqrt(dh)
    scores = scores.masked_fill(~valid, float("-inf"))

    # Analytic score VJP is useful for projection-only audits.  The CD path below
    # obtains the same quantities solely from paired softmax/value evaluations.
    if cfg.mode == "oracle":
        att = scores.softmax(-1)
        g_att = delta @ value.transpose(-2, -1)
        g_score_exact = att * (g_att - (att * g_att).sum(-1, keepdim=True))
        g_score_exact = g_score_exact.masked_fill(~valid, 0)

    n = int(cfg.n_probes)
    if n <= 0:
        raise ValueError("score-coordinate probing needs at least one probe")
    slots = min(n, T)
    # Independent random subsets for every sample/head/row.  Invalid positions
    # sort last, so the first min(n,m) entries are a uniform subset of the m-prefix.
    priorities = torch.rand((B, H, T, T), generator=gen, device=qkv.device, dtype=torch.float32)
    priorities.masked_fill_(~valid.view(1, 1, T, T), 2.0)
    chosen = priorities.argsort(dim=-1)[..., :slots]             # [B,H,T,slots]
    chosen = chosen.permute(3, 0, 1, 2).contiguous()             # [slots,B,H,T]
    slot_id = torch.arange(slots, device=qkv.device).view(slots, 1, 1, 1)
    active = slot_id < counts.view(1, 1, 1, T)
    active = active.expand(slots, B, H, T)
    weights = torch.where(counts <= n, torch.ones_like(counts, dtype=qkv.dtype),
                          counts.to(qkv.dtype) / n)               # [T]
    g_score = torch.zeros_like(scores)

    if cfg.mode == "oracle":
        picked = torch.gather(g_score_exact.unsqueeze(0).expand(slots, -1, -1, -1, -1),
                              -1, chosen.unsqueeze(-1)).squeeze(-1)
        contrib = picked * active.to(qkv.dtype) * weights.view(1, 1, 1, T)
        for j in range(slots):
            g_score.scatter_add_(-1, chosen[j].unsqueeze(-1), contrib[j].unsqueeze(-1))
        queries = 0.0
    else:
        eps = float(cfg.eps)  # directions are unit coordinates, so coord/norm conventions coincide
        if eps <= 0:
            raise ValueError("score-coordinate probing needs eps > 0")
        chunk = slots if cfg.probe_chunk <= 0 else max(1, min(int(cfg.probe_chunk), slots))
        for i0 in range(0, slots, chunk):
            ids = chosen[i0:i0 + chunk]
            act = active[i0:i0 + chunk]
            c = ids.shape[0]
            plus = scores.unsqueeze(0).expand(c, -1, -1, -1, -1).clone()
            minus = plus.clone()
            shift = act.to(qkv.dtype).unsqueeze(-1) * eps
            plus.scatter_add_(-1, ids.unsqueeze(-1), shift)
            minus.scatter_add_(-1, ids.unsqueeze(-1), -shift)
            y_plus = plus.softmax(-1) @ value.unsqueeze(0)
            y_minus = minus.softmax(-1) @ value.unsqueeze(0)
            if cfg.readout_noise > 0:
                y_plus = y_plus + cfg.readout_noise * torch.randn(
                    y_plus.shape, generator=gen, device=y_plus.device, dtype=y_plus.dtype)
                y_minus = y_minus + cfg.readout_noise * torch.randn(
                    y_minus.shape, generator=gen, device=y_minus.device, dtype=y_minus.dtype)
            deriv = ((y_plus - y_minus) * delta.unsqueeze(0)).sum(-1) / (2 * eps)
            contrib = deriv * act.to(qkv.dtype) * weights.view(1, 1, 1, T)
            for j in range(c):
                g_score.scatter_add_(-1, ids[j].unsqueeze(-1), contrib[j].unsqueeze(-1))
        queries = float(2 * n)

    _LAST_SCORE_VJP['g'] = g_score
    gk_heads = g_score.transpose(-2, -1) @ q / math.sqrt(dh)
    if dk == dh:
        gk = gk_heads.sum(1).transpose(1, 2)                      # shared K accumulates all heads
    else:
        gk = gk_heads.transpose(1, 2).reshape(B, T, int(d))
    return gk, queries


def make_h(f, v, n_batch_dims=1, local=None, readout_noise=0.0, gen=None):
    """Measurement function h(x_pert) = <v, f(x_pert)> per unit, tolerant to stacked batches.
    readout_noise > 0 adds i.i.d. Gaussian noise to every measured output coordinate of every query
    (a physical readout-noise model; v is applied digitally afterwards)."""
    B0 = v.shape[0]

    def h(xp):
        with torch.no_grad():
            y = f(xp)
            if readout_noise > 0:
                y = y + readout_noise * torch.randn(y.shape, generator=gen, device=y.device, dtype=y.dtype)
            k = y.shape[0] // B0
            vv = v.repeat(k, *([1] * (v.dim() - 1))) if k > 1 else v
            prod = y * vv
            if local is None:
                return prod.reshape(*prod.shape[:n_batch_dims], -1).sum(-1)
            if getattr(local, "kind", "2d") == "1d":
                return prod.sum(-1)                           # per output token
            return prod.sum(n_batch_dims)                     # keep spatial dims
    return h


def control_vjp(g, cfg, gen, n_batch_dims=1):
    """Control estimators.  'zero': no credit is propagated (downstream layers frozen).
    'noise': isotropic Gaussian noise with the per-sample norm that the oracle estimate with
    the same n would have (sqrt(1 + (d-1)/n) ||g||): same magnitude, zero information."""
    if cfg.mode == "zero":
        return torch.zeros_like(g), 0.0
    if cfg.mode == "noise":
        bs = g.shape[:n_batch_dims]
        gf = g.reshape(*bs, -1)
        d = gf.shape[-1]
        z = torch.randn(gf.shape, generator=gen, device=g.device, dtype=g.dtype)
        z = z / z.norm(dim=-1, keepdim=True).clamp_min(1e-30)
        target = gf.norm(dim=-1, keepdim=True) * math.sqrt(1 + (d - 1) / cfg.n_probes)
        return (z * target).reshape(g.shape), 0.0
    raise ValueError(cfg.mode)


def causal_score_coordinate_qk_vjp(qkv, v_out, cfg, gen, d, n_heads, dk=None, dv=None, window=None):
    """Same score-space probing as causal_score_coordinate_k_vjp, but returns the Q-side VJP as well
    (both follow from the score-VJP through the known bilinear map S = Q K^T / sqrt(d_h))."""
    B, T, _ = qkv.shape
    H = int(n_heads); dh = int(d) // H
    dk = int(d if dk is None else dk)
    gk, queries = causal_score_coordinate_k_vjp(qkv, v_out, cfg, gen, d, n_heads, dk, dv, window)
    # recover g_score from gk is not possible in general; recompute the score VJP estimate via the same
    # probes by calling the K routine on a transposed view is not exact either, so instead re-derive:
    # g_K[s] = sum_t g_S[t,s] q_t / sqrt(dh)  and  g_Q[t] = sum_s g_S[t,s] k_s / sqrt(dh).
    # We obtain g_S by solving nothing: the K routine exposes it through `_last_score_vjp`.
    g_score = _LAST_SCORE_VJP["g"]
    k0 = qkv[..., int(d):int(d) + dk]
    if dk == dh:
        k = k0.view(B, T, 1, dh).transpose(1, 2).expand(B, H, T, dh)
    else:
        k = k0.view(B, T, H, dh).transpose(1, 2)
    gq = (g_score @ k) / math.sqrt(dh)                              # [B,H,T,dh]
    gq = gq.transpose(1, 2).reshape(B, T, int(d))
    return gq, gk, queries


_LAST_SCORE_VJP = {"g": None}