"""Finite-domain control variates with explicit, auditable estimator contracts.

The returned estimate is unbiased for the sum of the current physical
contributions, conditional on the frozen predictor/proposal, when draws are
fresh, the oracle is correct, and sample count is fixed before drawing.
This is a classical control-variate identity, not a new unbiasedness theorem.
Neither clipping nor a subsequent nonlinear decoder preserves that identity.
"""
from dataclasses import dataclass
import numpy as np


@dataclass(frozen=True)
class FrozenControl:
    control: np.ndarray  # [receivers, terms, channels], contribution per term
    proposal: np.ndarray  # [receivers, terms], strictly positive probabilities
    integral: np.ndarray  # exact finite sum of the *same* frozen control


def freeze(control, proposal):
    h = np.array(control, dtype=np.float64, copy=True)
    q = np.array(proposal, dtype=np.float64, copy=True)
    if h.ndim != 3 or q.shape != h.shape[:2] or min(h.shape) < 1:
        raise ValueError("Expected nonempty control [P,K,C] and proposal [P,K]")
    if not np.isfinite(h).all() or not np.isfinite(q).all():
        raise ValueError("Control and proposal must be finite")
    if (q <= 0).any() or not np.allclose(q.sum(1), 1, rtol=0, atol=1e-12):
        raise ValueError("Proposal must have full support and rows summing to one")
    # Normalize within the accepted tolerance so the sampler and weights agree.
    q /= q.sum(1, keepdims=True)
    integral = h.sum(1)
    for arr in (h, q, integral):
        arr.setflags(write=False)
    return FrozenControl(h, q, integral)


def draw(snapshot, n, rng):
    if isinstance(n, bool) or not isinstance(n, (int, np.integer)) or n < 1:
        raise ValueError("n must be a positive integer fixed before sampling")
    u = rng.random((snapshot.proposal.shape[0], n))
    cdf = np.cumsum(snapshot.proposal, axis=1)
    cdf[:, -1] = 1.0
    return (u[..., None] >= cdf[:, None, :]).sum(-1)


def _indices(snapshot, indices):
    j = np.asarray(indices)
    if (j.ndim != 2 or j.shape[0] != len(snapshot.control) or j.shape[1] < 1
            or not np.issubdtype(j.dtype, np.integer)):
        raise ValueError("Indices must be a nonempty integer array [P,n]")
    if (j < 0).any() or (j >= snapshot.control.shape[1]).any():
        raise ValueError("Index outside finite physical domain")
    return j, np.arange(len(j))[:, None]


def correct(snapshot, indices, physical_values):
    """Correct one frozen prediction. Caller must respect the draw contract.

    Do not refit the control using these same values before this call. The
    lower-level API cannot detect correlations, biased oracles, or misuse of
    caller-supplied indices; use sample() to own the sampling boundary.
    """
    j, rows = _indices(snapshot, indices)
    f = np.asarray(physical_values, dtype=np.float64)
    if f.shape != j.shape + (snapshot.control.shape[2],) or not np.isfinite(f).all():
        raise ValueError("Physical values must be finite [P,n,C]")
    residual = (f - snapshot.control[rows, j]) / snapshot.proposal[rows, j, None]
    return snapshot.integral + residual.mean(1)


def sample(snapshot, oracle, n, rng):
    """Draw fresh indices, call oracle(indices)->[P,n,C], then correct.

    Returns estimate, indices, values. Commit evidence to memory only after
    this returns; the predictor is frozen independently of these samples.
    """
    j = draw(snapshot, n, rng)
    f = np.asarray(oracle(j), dtype=np.float64)
    return correct(snapshot, j, f), j, f


def exact_mse(physical_table, snapshot, n=1):
    """Audit-only conditional MSE, averaged over channels, for fixed truth.

    This enumerates every physical term. Never expose it to the online policy
    when claiming a sparse-ray budget. Useful for deterministic validation.
    """
    f = np.asarray(physical_table, dtype=np.float64)
    if f.shape != snapshot.control.shape or not np.isfinite(f).all() or n < 1:
        raise ValueError("Invalid physical audit table or sample count")
    residual = f - snapshot.control
    second = (residual**2 / snapshot.proposal[..., None]).sum(1)
    squared_mean = residual.sum(1)**2
    return np.maximum(second - squared_mean, 0).mean(-1) / n


def residual_metric(coefficients, proposal, n=1, channel_metric=None):
    """G=(diag(c_j^T Q c_j/q_j)-C^T Q C)/n.

    For latent term vector x, contribution j is c_j*x_j. If the control uses
    its true conditional mean and covariance is P, posterior-averaged
    corrected-estimator risk is trace(G P). This is a *model-dependent* value
    calculation; core estimator unbiasedness does not require that model.
    """
    c = np.asarray(coefficients, dtype=float)
    q = np.asarray(proposal, dtype=float)
    if c.ndim != 2 or q.shape != (c.shape[1],) or n < 1 or (q <= 0).any():
        raise ValueError("Expected C [channels,K], positive q [K], and n>=1")
    if not np.isfinite(c).all() or not np.isfinite(q).all() or not np.isclose(q.sum(), 1):
        raise ValueError("Invalid finite coefficients/probability sum")
    weight = np.eye(c.shape[0]) if channel_metric is None else np.asarray(channel_metric, float)
    if (weight.shape != (c.shape[0], c.shape[0]) or not np.allclose(weight, weight.T)
            or not np.isfinite(weight).all() or np.linalg.eigvalsh(weight).min() < -1e-12):
        raise ValueError("Channel metric must be symmetric positive semidefinite")
    gram = c.T @ weight @ c
    return (np.diag(np.diag(gram)/q) - gram) / n


def hoeffding_radius(snapshot, upper_bounds, n, delta=0.05):
    """Per-receiver simultaneous-channel fixed-n bound; NOT an anytime bound.

    Requires actual physical 0<=f_jc<=upper_bounds_jc, fixed snapshot and iid
    categorical samples. It is generally very conservative at low ray count.
    For a whole-frame statement use delta/number_of_receivers.
    """
    b = np.asarray(upper_bounds, float)
    if b.shape != snapshot.control.shape or not np.isfinite(b).all() or (b < 0).any():
        raise ValueError("Need finite nonnegative physical bounds [P,K,C]")
    if n < 1 or not 0 < delta < 1:
        raise ValueError("Need n>=1 and 0<delta<1")
    low = (-snapshot.control / snapshot.proposal[..., None]).min(1)
    high = ((b - snapshot.control) / snapshot.proposal[..., None]).max(1)
    channels = b.shape[2]
    return (high-low) * np.sqrt(np.log(2*channels/delta)/(2*n))


def proposal_from_bound(bound, visibility, trusted=None, active=False, floor=0.1):
    """Full-support heuristic; optimality and calibration are not promised."""
    b = np.asarray(bound, float)
    p = np.asarray(visibility, float)
    if b.ndim != 3 or p.shape != b.shape[:2] or not 0 < floor <= 1:
        raise ValueError("Invalid proposal inputs")
    if not np.isfinite(b).all() or (b < 0).any() or not np.isfinite(p).all() or ((p < 0)|(p > 1)).any():
        raise ValueError("Bound must be nonnegative and visibility in [0,1]")
    score = np.linalg.norm(b, axis=-1)
    if active:
        uncertainty = p * (1-p) + 0.04
        if trusted is not None:
            trust = np.asarray(trusted, dtype=bool)
            if trust.shape != p.shape:
                raise ValueError("Trust mask shape mismatch")
            uncertainty = np.where(trust, 0.0, uncertainty)
        score *= np.sqrt(uncertainty)
    total = score.sum(1, keepdims=True)
    base = np.divide(score, total, out=np.full_like(score, 1/score.shape[1]), where=total > 0)
    return (1-floor)*base + floor/score.shape[1]