import json
import numpy as np


def pinv_truncated(C, rel_tol=1e-10):
    u, s, vt = np.linalg.svd(C, full_matrices=False)
    if len(s) == 0 or s[0] == 0:
        return np.zeros((C.shape[1], C.shape[0]))
    keep = s > rel_tol * s[0]
    return (vt[keep].T / s[keep]) @ u[:, keep].T


def cur_factors(W, rows, cols, rel_tol=1e-8):
    rows, cols = np.asarray(rows, dtype=int), np.asarray(cols, dtype=int)
    A = W[:, cols]
    B = W[rows, :]
    M = pinv_truncated(W[np.ix_(rows, cols)], rel_tol)
    return A, M, B


def reconstruct(factors):
    A, M, B = factors
    return A @ M @ B


def _top_with_retention(scores, old, k, retain):
    old = np.asarray(old, dtype=int)
    nold = max(1, int(round(retain * k)))
    retained = list(old[np.argsort(scores[old])[::-1][:nold]])
    chosen = retained[:]
    for idx in np.argsort(scores)[::-1]:
        if idx not in chosen:
            chosen.append(int(idx))
        if len(chosen) == k:
            break
    return np.asarray(chosen, dtype=int)


def residual_leverage_refresh(W, old_rows, old_cols, probes=8, retain=0.25,
                              rng=None, rel_tol=1e-8):
    rng = np.random.default_rng() if rng is None else rng
    A, M, B = cur_factors(W, old_rows, old_cols, rel_tol)
    R = W - A @ M @ B
    # Row scores from RP and column scores from P^T R.
    Pcol = rng.standard_normal((W.shape[1], probes))
    Qr, _ = np.linalg.qr(R @ Pcol, mode='reduced')
    Prow = rng.standard_normal((W.shape[0], probes))
    Qc, _ = np.linalg.qr(R.T @ Prow, mode='reduced')
    row_scores = np.sum(Qr * Qr, axis=1)
    col_scores = np.sum(Qc * Qc, axis=1)
    rows = _top_with_retention(row_scores, old_rows, len(old_rows), retain)
    cols = _top_with_retention(col_scores, old_cols, len(old_cols), retain)
    return cur_factors(W, rows, cols, rel_tol), rows, cols, np.linalg.norm(R) / np.linalg.norm(W)


def rel_error(W, factors):
    return np.linalg.norm(W - reconstruct(factors)) / np.linalg.norm(W)


def factor_param_count(W, factors):
    A, M, B = factors
    return A.size + M.size + B.size


def run(seed=7):
    rng = np.random.default_rng(seed)
    m, n, r, k = 64, 80, 8, 12
    # A changing low-rank-plus-noise matrix: adaptive residual scores should track new rows/cols.
    U = rng.standard_normal((m, r)); U, _ = np.linalg.qr(U)
    V = rng.standard_normal((n, r)); V, _ = np.linalg.qr(V)
    singular = np.linspace(3.0, 0.5, r)
    W0 = (U * singular) @ V.T
    initial_rows = np.arange(k)
    initial_cols = np.arange(k)
    fixed = cur_factors(W0, initial_rows, initial_cols)
    random = cur_factors(W0, rng.choice(m, k, replace=False), rng.choice(n, k, replace=False))
    exact = rel_error(W0, cur_factors(W0, np.arange(m), np.arange(n)))
    initial_err = rel_error(W0, fixed)
    random_err = rel_error(W0, random)

    # Perturb a previously unselected block, then compare stale, random refresh, adaptive refresh.
    W1 = W0.copy()
    new_rows = np.arange(40, 52); new_cols = np.arange(55, 67)
    W1[np.ix_(new_rows, new_cols)] += 0.8 * rng.standard_normal((len(new_rows), len(new_cols)))
    stale_err = rel_error(W1, fixed)
    _, rr, cc, _ = residual_leverage_refresh(W1, initial_rows, initial_cols, probes=16, retain=.25, rng=np.random.default_rng(seed+1))
    adaptive = cur_factors(W1, rr, cc)
    adaptive_err = rel_error(W1, adaptive)
    random2 = cur_factors(W1, rng.choice(m, k, replace=False), rng.choice(n, k, replace=False))
    random2_err = rel_error(W1, random2)
    # Formula/evaluation equivalence on a batch.
    X = rng.standard_normal((11, m))
    y1 = X @ reconstruct(adaptive)
    A, M, B = adaptive
    y2 = ((X @ A) @ M) @ B
    formula_gap = np.max(np.abs(y1-y2))
    return dict(seed=seed, shape=[m,n], rank=r, cross_rank=k,
                exact_dense_error=exact, initial_error=initial_err,
                random_initial_error=random_err, stale_after_change=stale_err,
                adaptive_after_change=adaptive_err, random_after_change=random2_err,
                adaptive_rows=rr.tolist(), adaptive_cols=cc.tolist(),
                formula_max_abs_gap=float(formula_gap), dense_params=m*n,
                cur_params=factor_param_count(W1, adaptive), compression=(m*n)/factor_param_count(W1, adaptive))


if __name__ == '__main__':
    print(json.dumps(run(), indent=2))