import json, time
import numpy as np
from active_set_router import active_set_simplex, projected_gradient, objective

np.set_printoptions(precision=4, suppress=True)
rng = np.random.default_rng(1263)
E = 16
# Orthogonal feature construction makes the exact Gram spectrum known while
# retaining matrix-free Z/Z.T products in the solver.
Q, _ = np.linalg.qr(rng.normal(size=(E,E)))
U, _ = np.linalg.qr(rng.normal(size=(E,E)))
xtrue = np.zeros(E); xtrue[[0,3,7,12]] = [0.05, 0.15, 0.30, 0.50]
rows=[]
for lam in [1e-4, 1e-3, 1e-2, 1e-1, 1.0]:
    # eigenvalues of Z'Z range from 1 to 1e4; ridge controls kappa
    eig = np.geomspace(1.0, 1e4, E)
    Z = Q @ np.diag(np.sqrt(eig)) @ U.T
    A = Z.T@Z + lam*np.eye(E)
    b = A@xtrue
    cond = np.linalg.eigvalsh(A)[-1]/np.linalg.eigvalsh(A)[0]
    t0=time.perf_counter(); xa, ma = active_set_simplex(Z,b,lam=lam,cg_tol=1e-11,max_pivots=300); ta=time.perf_counter()-t0
    t0=time.perf_counter(); xp, pg = projected_gradient(Z,b,lam=lam,steps=100000,tol=1e-10); tp=time.perf_counter()-t0
    # KKT residual after choosing equality multiplier from free coordinates
    g0 = A@xa-b
    free = xa > 1e-8
    nu = -np.mean(g0[free]) if np.any(free) else 0.0
    station = np.max(np.abs((g0+nu)[free])) if np.any(free) else np.nan
    dual = np.min((g0+nu)[~free]) if np.any(~free) else 0.0
    rows.append(dict(lambda_=lam, condition=cond, sqrt_condition=np.sqrt(cond),
                     active_cg_matvecs=ma['cg_matvecs'], pivots=ma['pivots'],
                     active_free=ma['free'], pg_steps=pg,
                     active_objective=objective(Z,b,xa,lam), pg_objective=objective(Z,b,xp,lam),
                     simplex_error=float(abs(xa.sum()-1)), min_x=float(xa.min()),
                     kkt_free=float(station), kkt_bound_min=float(dual),
                     active_time=ta, pg_time=tp))
# Separate spectrum sweep: predicted CG dependence on sqrt(kappa), using the
# same simplex problem but varying the raw Gram dynamic range.
scaling=[]
lam=1e-3
for ratio in [1, 10, 100, 1000, 10000]:
    eig=np.geomspace(1., ratio, E)
    Z=Q@np.diag(np.sqrt(eig))@U.T
    A=Z.T@Z+lam*np.eye(E); b=A@xtrue
    cond=np.linalg.eigvalsh(A)[-1]/np.linalg.eigvalsh(A)[0]
    xa,ma=active_set_simplex(Z,b,lam=lam,cg_tol=1e-11,max_pivots=300)
    xp,pg=projected_gradient(Z,b,lam=lam,steps=100000,tol=1e-10)
    scaling.append(dict(raw_ratio=ratio, condition=cond, sqrt_condition=np.sqrt(cond),
                        cg_matvecs=ma['cg_matvecs'], pg_steps=pg,
                        objective_gap=objective(Z,b,xp,lam)-objective(Z,b,xa,lam)))
result={'ridge_sweep':rows,'spectrum_sweep':scaling,
        'predictions':[
          'Increasing ridge decreases kappa(A) monotonically and should reduce CG work.',
          'Across spectra, CG work should track sqrt(kappa) more closely than projected-gradient steps track kappa.',
          'The active-set output should have exact simplex feasibility and nonnegative coefficients, with sparse support.'
        ]}
with open('results.json','w') as f: json.dump(result,f,indent=2)
print(json.dumps(result,indent=2))