"""Exact weighted linear-probe coreset by preserving normal equations."""
import numpy as np


def numerical_rank(x, tol=1e-6):
    s = np.linalg.svd(x, compute_uv=False)
    return int(np.sum(s > (tol * s[0] if len(s) and s[0] else 0)))


def fit_min_norm(x, y, weights=None):
    if weights is None:
        weights = np.ones(len(x))
    a = x.T @ (weights[:, None] * x)
    b = y.T @ (weights[:, None] * x)
    return b @ np.linalg.pinv(a)


def residual_coreset(x, y, tol=1e-10, max_iter=None):
    """Return <= (m+1)r examples preserving the full optimum's normal equations.

    For W*, each atom is residual_i outer x_i.  A nonnegative nullspace move
    among k+1 atoms preserves sum w_i residual_i outer x_i, where k=m*d.
    Local elimination avoids large full nullspace SVDs.
    """
    n, d = x.shape
    m = y.shape[1]
    weights = np.ones(n)
    W = fit_min_norm(x, y)
    r = numerical_rank(x)
    R = ((y - x @ W.T)[:, :, None] * x[:, None, :]).reshape(n, m*d)
    active = list(range(n))
    k = R.shape[1]
    target = (m+1) * max(r, 1)
    if max_iter is None: max_iter = n + 5
    it = 0
    while len(active) > target and it < max_iter:
        # Any k+1 atoms are linearly dependent. If this particular block has
        # rank deficiency, its null vector still gives a valid elimination.
        block = np.asarray(active[:min(k+1, len(active))], dtype=int)
        _, s, vh = np.linalg.svd(R[block].T, full_matrices=True)
        c = vh[-1]
        if np.linalg.norm(c) < tol:
            break
        if np.all(c <= tol): c = -c
        pos = c > tol
        if not np.any(pos): break
        wb = weights[block]
        t = np.min(wb[pos] / c[pos])
        weights[block] -= t*c
        weights[np.abs(weights) < 1e-12] = 0.0
        active = [i for i in active if weights[i] > 1e-12]
        it += 1
    active = np.asarray([i for i in active if weights[i] > 1e-12], dtype=int)
    return active, weights[active], W, r, it


def weighted_loss(x, y, idx, weights, W):
    e = x[idx] @ W.T - y[idx]
    return float(np.sum(weights[:, None] * e*e))