from dataclasses import dataclass, asdict, replace


@dataclass
class ZBPConfig:
    """Configuration of the zeroth-order VJP estimator used inside every ZBPBlock.

    mode:
        'bp'           exact reverse-mode autograd through the block (baseline)
        'zero'         control: no activation error is propagated (upstream layers get no credit)
        'noise'        control: random error with the norm an n-probe oracle estimate would have
        'fa'           feedback alignment: fixed random linear map of the incoming error (skip path exact)
        'dfa'          direct feedback alignment: fixed random projection of the output error (no chain)
        'exact'        same as 'bp' but routed through the ZBP backward (debug/consistency check)
        'oracle'       Stage A: exact g = J^T v computed internally, only its random
                       projection (1/n) sum_i u_i (u_i^T g) is returned  (projection noise only)
        'forward'      one-sided finite difference  (h(x+eps u) - h(x)) / eps
        'cd'           Stage B: central difference   (h(x+eps u) - h(x-eps u)) / (2 eps)
        'richardson'   Richardson-corrected central difference  (4 D_{eps/2} - D_eps) / 3
        'ml'           Stage C: randomized multilevel debiased central difference (unbiased)
        'ml_richardson'  randomized multilevel built on the Richardson sequence (unbiased)
    n_probes:  number of random directions per sample unit per backward call.
    probe:     'rademacher' | 'gaussian' | 'orthogonal' | 'hadamard' | 'coordinate'.
               All families are normalized so that E[u u^T] = I_d (entries are O(1)).
    eps:       finite-difference step.  With eps_mode='coord' the perturbation is x + eps*u
               (each coordinate moves by ~eps, Euclidean norm eps*sqrt(d)); with
               eps_mode='norm' it is x + eps*u/sqrt(d) (fixed Euclidean norm eps).
    alpha:     survival exponent of the multilevel truncation level, P(N >= k) = 2^{-alpha k}.
    max_level: hard cap on the multilevel index (P(N > max_level) = 2^{-alpha (max_level+1)}).
    level_sampling: 'per_sample' (independent N for every (probe, sample) pair — the physical
               cost model) or 'per_probe' (one N per probe shared by the batch — cheaper to simulate).
    batch_probes: evaluate all probes of a backward call in one batched forward (simulation speed).
    """
    mode: str = "oracle"
    n_probes: int = 4
    probe: str = "rademacher"
    eps: float = 1e-2
    eps_mode: str = "coord"
    alpha: float = 2.0
    max_level: int = 12
    level_sampling: str = "per_sample"
    batch_probes: bool = True
    probe_chunk: int = 0        # >0: evaluate at most this many probes per batched forward (bounds memory)
    param_mode: str = "local"   # 'local': J_theta^T v via in-block autograd from the incoming error
    readout_noise: float = 0.0  # std of additive Gaussian noise on every measured output coordinate (per query)
    eps_min: float = 0.0        # multilevel ladder is truncated at eps_k >= eps_min (residual bias O(eps_min^2))
    surrogate: bool = False     # learned linear control variate J^T ~ A per block (fitted online from the ZO estimates)
    surrogate_ema: float = 0.9  # EMA factor of the regression statistics E[g v^T], E[v v^T]
    surrogate_ridge: float = 1e-3

    def replace(self, **kw):
        return replace(self, **kw)

    def asdict(self):
        return asdict(self)

    @property
    def is_zo(self):
        return self.mode in ("forward", "cd", "richardson", "ml", "ml_richardson")

    def expected_queries_per_probe(self):
        """Expected number of block evaluations (per sample unit) per probe direction."""
        if self.mode == "oracle" or self.mode in ("bp", "exact"):
            return 0.0
        if self.mode == "forward":
            return 1.0  # plus one shared base evaluation per backward call
        if self.mode == "cd":
            return 2.0
        if self.mode == "richardson":
            return 4.0
        q = 2.0 ** (-self.alpha)
        # E[N+1] = sum_{k>=0} Q_k = 1/(1-q)  (ignoring the cap)
        en1 = 1.0 / (1.0 - q)
        if self.mode == "ml":
            return 2.0 * en1
        if self.mode == "ml_richardson":
            return 2.0 * (en1 + 1.0)  # levels 0..N+1
        raise ValueError(self.mode)