import json
import math
import numpy as np

SEED = 1515

def invsqrt(A, eps=1e-12):
    w, V = np.linalg.eigh((A + A.T) * 0.5)
    w = np.maximum(w, eps)
    return (V * (1.0 / np.sqrt(w))) @ V.T

def metric_from_history(G, beta, lam, rho, alpha):
    p, m = G.shape
    S = beta * np.zeros((p, p)) + (1-beta) * (G @ G.T) / m + lam * np.eye(p)
    R = invsqrt(S)
    if alpha == 1.0:
        # Projection is mathematically unused in this baseline branch.
        P = np.zeros((p, p))
        Q = R @ R
    else:
        K = G.T @ np.linalg.solve(S, G) + rho * np.eye(m)
        P = R @ G @ np.linalg.solve(K, G.T) @ R
        Q = R @ ((1-alpha) * P + alpha * np.eye(p)) @ R
    return S, P, Q

def predicted_threshold(H, Q):
    # For general noncommuting H,Q this is the exact quadratic GD threshold.
    A = (Q @ H + H @ Q) * 0.5
    return 2.0 / np.max(np.linalg.eigvalsh(A))

def run_fixed(H, G, eta, beta, lam, rho, alpha, steps=100):
    _, _, Q = metric_from_history(G, beta, lam, rho, alpha)
    theta = np.array([1.0, 1.0])
    losses = []
    for _ in range(steps):
        losses.append(0.5 * theta @ H @ theta)
        theta = theta - eta * Q @ (H @ theta)
        if not np.all(np.isfinite(theta)) or np.linalg.norm(theta) > 1e8:
            return False, losses
    return losses[-1] < losses[0] and np.linalg.norm(theta) < 1.0, losses

def observed_threshold(H, G, beta, lam, rho, alpha):
    pred = predicted_threshold(H, metric_from_history(G,beta,lam,rho,alpha)[2])
    # binary search for largest tested stable eta; stability means monotone decay over 300 steps
    lo, hi = 0.0, pred * 1.8
    for _ in range(30):
        mid = (lo + hi) / 2
        ok, _ = run_fixed(H,G,mid,beta,lam,rho,alpha,steps=300)
        if ok: lo = mid
        else: hi = mid
    return lo, pred

def quadratic_training(H, Gbase, eta, lam, rho, alpha, rng, steps=160, noise=0.25):
    p = H.shape[0]
    theta = np.array([1.0, 1.0])
    hist = []
    losses = []
    for t in range(steps):
        g = H @ theta + noise * rng.normal(size=p)
        hist.append(g.copy())
        G = np.column_stack(hist[-Gbase:])
        _, _, Q = metric_from_history(G, 0.0, lam, rho, alpha)
        theta -= eta * Q @ g
        losses.append(0.5 * theta @ H @ theta)
    return float(np.mean(losses[-30:])), float(np.mean(losses)), losses

def main():
    H = np.diag([12.0, 1.0])
    # The history only identifies the stiff coordinate: this makes the weak-support claim explicit.
    G = np.array([[1.0, 0.8, 1.2, 0.9], [0.0, 0.0, 0.0, 0.0]])
    beta, rho, alpha = 0.0, 0.15, 0.08

    # Prediction 1: increasing ridge shrinks the supported metric eigenvalue and raises the stability limit.
    ridge_rows = []
    for lam in [0.01, 0.03, 0.1, 0.3, 1.0]:
        obs, pred = observed_threshold(H,G,beta,lam,rho,alpha)
        Q = metric_from_history(G,beta,lam,rho,alpha)[2]
        ridge_rows.append({'lambda':lam, 'q_supported':float(Q[1,1]), 'pred_eta_c':pred, 'observed_eta_c':obs})

    # Prediction 2: increasing rho weakens projection-supported directions and raises eta_c.
    rho_rows = []
    for r in [0.01, 0.03, 0.1, 0.3, 1.0]:
        obs, pred = observed_threshold(H,G,beta,0.1,r,alpha)
        Q = metric_from_history(G,beta,0.1,r,alpha)[2]
        rho_rows.append({'rho':r, 'q_supported':float(Q[0,0]), 'pred_eta_c':pred, 'observed_eta_c':obs})

    # Prediction 3: alpha controls unsupported directions: at alpha=0 they receive zero metric.
    alpha_rows = []
    for a in [0.0, 0.02, 0.1, 0.3, 1.0]:
        Q = metric_from_history(G,beta,0.1,rho,a)[2]
        alpha_rows.append({'alpha':a, 'q_unsupported':float(Q[1,1]), 'q_supported':float(Q[0,0])})

    # Secondary matched-step comparison on noisy quadratic, averaging fixed seeds.
    baseline, idea = [], []
    eta = 0.055
    for seed in range(8):
        rb = np.random.default_rng(SEED + seed)
        ri = np.random.default_rng(SEED + seed)
        baseline.append(quadratic_training(H, 4, eta, 1e-12, 0.0, 1.0, rb)[0])
        idea.append(quadratic_training(H, 4, eta, 0.1, rho, alpha, ri)[0])
    result = {
        'settings': {'H_diag':[12.0,1.0], 'history_rank':1, 'beta':beta, 'rho':rho, 'alpha':alpha},
        'ridge_sweep':ridge_rows, 'rho_sweep':rho_rows, 'alpha_sweep':alpha_rows,
        'comparison': {'baseline_mean_final_loss':float(np.mean(baseline)), 'baseline_std':float(np.std(baseline)), 'idea_mean_final_loss':float(np.mean(idea)), 'idea_std':float(np.std(idea)), 'eta':eta},
        'checks': {
            'ridge_predicted_monotonic': bool(all(ridge_rows[i]['pred_eta_c'] < ridge_rows[i+1]['pred_eta_c'] for i in range(4))),
            'ridge_observed_monotonic': bool(all(ridge_rows[i]['observed_eta_c'] < ridge_rows[i+1]['observed_eta_c'] for i in range(4))),
            'rho_predicted_monotonic': bool(all(rho_rows[i]['pred_eta_c'] < rho_rows[i+1]['pred_eta_c'] for i in range(4))),
            'rho_observed_monotonic': bool(all(rho_rows[i]['observed_eta_c'] < rho_rows[i+1]['observed_eta_c'] for i in range(4))),
            'unsupported_zero_at_alpha0': bool(alpha_rows[0]['q_unsupported'] < 1e-10),
            'unsupported_increases_with_alpha': bool(all(alpha_rows[i]['q_unsupported'] < alpha_rows[i+1]['q_unsupported'] for i in range(4)))
        }
    }
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result, indent=2))

if __name__ == '__main__': main()