import json
import numpy as np
from adaptive_cur import cur_factors, residual_leverage_refresh, rel_error


def one(seed):
    rng = np.random.default_rng(seed)
    m, n, r, k = 64, 80, 8, 12
    U, _ = np.linalg.qr(rng.standard_normal((m, r)))
    V, _ = np.linalg.qr(rng.standard_normal((n, r)))
    W0 = (U * np.linspace(3.0, .5, r)) @ V.T
    rows, cols = np.arange(k), np.arange(k)
    stale = cur_factors(W0, rows, cols)
    W1 = W0.copy()
    # New structure is deliberately outside the original index sets.
    ir = np.arange(40, 52); ic = np.arange(55, 67)
    W1[np.ix_(ir, ic)] += .8 * rng.standard_normal((len(ir), len(ic)))
    stale_e = rel_error(W1, stale)
    af, ar, ac, _ = residual_leverage_refresh(W1, rows, cols, probes=16, retain=.25,
                                               rng=np.random.default_rng(seed + 100))
    adaptive_e = rel_error(W1, af)
    rf = cur_factors(W1, rng.choice(m, k, replace=False), rng.choice(n, k, replace=False))
    random_e = rel_error(W1, rf)
    return stale_e, adaptive_e, random_e

vals = np.array([one(s) for s in range(20)])
print(json.dumps({
    'seeds': 20,
    'mean_relative_errors': dict(stale=float(vals[:,0].mean()), adaptive=float(vals[:,1].mean()), random=float(vals[:,2].mean())),
    'median_relative_errors': dict(stale=float(np.median(vals[:,0])), adaptive=float(np.median(vals[:,1])), random=float(np.median(vals[:,2]))),
    'adaptive_beats_stale': int(np.sum(vals[:,1] < vals[:,0])),
    'adaptive_beats_random': int(np.sum(vals[:,1] < vals[:,2])),
    'per_seed': vals.tolist()
}, indent=2))