"""Sequential residual elimination with fixed budget and exact partial sums.

Do not supply unverified entries as known. Zero proposal support is legal
only when the contribution is exact. This module does not validate geometry.
"""
import numpy as np


def prepare(bound, prior, values, known, score):
    b, p, v = np.asarray(bound, float), np.asarray(prior, float), np.asarray(values, float)
    k, s = np.asarray(known, bool), np.asarray(score, float)
    if (b.ndim != 3 or p.shape != b.shape[:2] or v.shape != p.shape or k.shape != p.shape
            or s.shape != p.shape or not all(np.isfinite(x).all() for x in (b,p,v,s))
            or (b < 0).any() or ((p < 0)|(p > 1)).any() or ((v < 0)|(v > 1)).any() or (s <= 0).any()):
        raise ValueError("Invalid finite control or strictly positive base scores")
    h = b*np.where(k, v, p)[..., None]
    scores = np.where(k, 0, s)
    total = scores.sum(1, keepdims=True)
    q = np.divide(scores, total, out=np.zeros_like(scores), where=total > 0)
    return h.copy(), q


def eliminate(h, q, oracle, n, rng):
    """oracle(rows,j) supplies *current exact* RGB for one term per row.

Each selected term is assimilated after its martingale estimate is formed.
No duplicates; no queries for rows already complete. Exact completion is
determined by the initial support size, not by observed sample values.
"""
    h, q = np.array(h, float, copy=True), np.array(q, float, copy=True)
    if (h.ndim != 3 or q.shape != h.shape[:2] or not np.isfinite(h).all()
            or not np.isfinite(q).all() or (q < 0).any()
            or not np.all(np.isclose(q.sum(1), 1, atol=1e-12) | (q.sum(1) == 0))
            or isinstance(n, bool) or not isinstance(n, (int, np.integer)) or n < 1):
        raise ValueError("Finite control, normalized nonnegative support and fixed positive budget required")
    initial_count = (q > 0).sum(1)
    scores = q.copy()
    estimates = np.zeros((len(h), h.shape[2]))
    selections = []
    for _ in range(n):
        integral = h.sum(1)
        active = np.flatnonzero(scores.sum(1) > 0)
        y = integral.copy()
        if len(active):
            weights = scores[active]/scores[active].sum(1, keepdims=True)
            cdf = np.minimum(np.cumsum(weights, 1), 1.0)
            last = weights.shape[1]-1-np.argmax(weights[:, ::-1] > 0, axis=1)
            # Floating-point summation must not give a trailing zero-support
            # bin a tiny spurious interval near one.
            cdf[np.arange(weights.shape[1])[None, :] >= last[:, None]] = 1.0
            j = (rng.random(len(active))[:, None] >= cdf).sum(1)
            f = np.asarray(oracle(active, j), float)
            if f.shape != (len(active), h.shape[2]) or not np.isfinite(f).all():
                raise ValueError("Oracle returned invalid physical values")
            y[active] += (f-h[active, j])/weights[np.arange(len(active)), j, None]
            h[active, j] = f
            scores[active, j] = 0
            selections.append((active.copy(), j.copy()))
        estimates += y/n
    complete = initial_count <= n
    estimates[complete] = h.sum(1)[complete]
    return estimates, selections


def exact_risk_two(truth, h, q):
    """Audit-only expected MSE for eliminate(..., n=2). No online access.

Works only when zero-support terms really are exact; otherwise raises.
"""
    f, h, q = np.asarray(truth), np.asarray(h), np.asarray(q)
    r = f-h
    if np.max(np.abs(np.where((q == 0)[..., None], r, 0)), initial=0) > 1e-10:
        raise ValueError("False certificate: nonzero residual has zero support")
    s = np.divide(r*r, q[..., None], out=np.zeros_like(r), where=q[..., None] > 0).sum(1)
    total = r.sum(1)
    v1 = s-total*total
    v2 = s*(1-(q*q).sum(1))[:, None]-(r*r).sum(1)-total*total+2*total*(q[..., None]*r).sum(1)
    result = np.maximum((v1+v2).mean(-1)/4, 0)
    result[(q > 0).sum(1) <= 2] = 0
    return result


def covariance(residual, proposal):
    """One-sample covariance on positive support; used in proof witnesses."""
    r, q = np.asarray(residual, float), np.asarray(proposal, float)
    use = q > 0
    if np.any(np.abs(r[~use]) > 1e-12):
        raise ValueError("Uncertified zero support")
    return (r[use].T/q[use])@r[use]-np.outer(r.sum(0), r.sum(0))