import json
import numpy as np
from pathlib import Path


def controller(gamma, H, S, H0, S0, rho=0.10, gmin=1e-7, gmax=0.05, eps0=1e-6):
    score = (H-H0)/max(H0, eps0) + (S-S0)/max(S0, eps0)
    return float(np.clip(gamma*np.exp(rho*score), gmin, gmax)), float(score)


def ridge_fit(X, y, gamma):
    n, d = X.shape
    A = X.T @ X / n + gamma*np.eye(d)
    return np.linalg.solve(A, X.T @ y / n)


def diagnostics(X, eta=0.30, eps=1e-10):
    n = X.shape[0]
    lam = np.linalg.eigvalsh(X.T @ X / n + eps*np.eye(X.shape[1]))
    inv = 1.0/lam
    H = float(np.mean(inv))
    M = float(np.mean(np.where(lam >= eta, inv, 0.0)))
    S = float(np.mean(np.where(lam < eta, inv, 0.0)))
    return H, M, S, float(lam.min()), lam


def make_data(rng, n, d, hard_edge=False):
    # A controlled spectrum: only the final coordinates become nearly unobserved.
    eig = np.ones(d)
    if hard_edge:
        eig[-d//4:] = np.geomspace(1e-4, 2e-2, d//4)
    else:
        eig[-d//4:] = np.geomspace(0.35, 1.0, d//4)
    X = rng.normal(size=(n,d)) * np.sqrt(eig)[None,:]
    beta = rng.normal(size=d) / np.sqrt(d)
    y = X @ beta + 0.10*rng.normal(size=n)
    return X, y, beta, eig


def verify_math():
    rng = np.random.default_rng(123)
    X, _, _, _ = make_data(rng, 128, 24, hard_edge=True)
    H, M, S, _, lam = diagnostics(X, eta=0.30)
    split_error = abs(H-(M+S))
    # The closed-form ridge solution is also checked against its normal equations.
    w = ridge_fit(X, rng.normal(size=128), 0.07)
    normal_eq = np.linalg.norm((X.T@X/128 + .07*np.eye(24))@w - X.T@rng.normal(size=128))
    # use a fresh explicit RHS for the actual residual check
    yy = rng.normal(size=128); ww = ridge_fit(X, yy, .07)
    residual = np.linalg.norm((X.T@X/128 + .07*np.eye(24))@ww - X.T@yy/128)
    return {"H":H, "M":M, "S":S, "decomposition_abs_error":split_error,
            "min_eigenvalue":float(lam.min()), "normal_equation_residual":float(residual)}


def run():
    rng = np.random.default_rng(2025)
    n, d, ntest, reps = 192, 48, 512, 80
    # Calibration represents the healthy-spectrum regime, as prescribed by the controller.
    cal_H, cal_S = [], []
    for _ in range(20):
        X, _, _, _ = make_data(rng, n, d, hard_edge=False)
        h, _, s, _, _ = diagnostics(X, eta=0.30)
        cal_H.append(h); cal_S.append(s)
    H0, S0 = float(np.median(cal_H)), float(np.median(cal_S))
    rows = {k: [] for k in ["ridgeless","fixed_ridge","adaptive"]}
    gammas, scores, hvals = [], [], []
    for hard in [False, True]:
        for _ in range(reps):
            X, y, beta, _ = make_data(rng, n, d, hard_edge=hard)
            H, M, S, mineig, _ = diagnostics(X, eta=0.30)
            # Initialize weakly; update once from this feature batch.
            ga, score = controller(1e-4, H, S, H0, S0)
            for name, gamma in [("ridgeless",0.0),("fixed_ridge",0.01),("adaptive",ga)]:
                w = ridge_fit(X, y, gamma)
                # Test corruption is additive feature noise, exposing unstable directions.
                Xt, _, _, _ = make_data(rng, ntest, d, hard_edge=hard)
                yt = Xt @ beta + 0.10*rng.normal(size=ntest)
                pred = Xt @ w
                losses = (pred-yt)**2
                rows[name].append({"hard":hard, "mse":float(losses.mean()),
                                   "p99":float(np.quantile(losses,.99)),
                                   "max_output":float(np.max(np.abs(pred))),
                                   "mineig":mineig})
            if hard: gammas.append(ga); scores.append(score); hvals.append(H)
    summary = {}
    for name, vals in rows.items():
        summary[name] = {}
        for regime, label in [(False,"healthy"),(True,"hard_edge")]:
            a = [v for v in vals if v["hard"] == regime]
            summary[name][label] = {k:float(np.mean([x[k] for x in a])) for k in ["mse","p99","max_output","mineig"]}
    # Robustness signal is the hard-edge tail loss; report paired regime counts and controller behavior.
    return {"math":verify_math(), "calibration":{"H0":H0,"S0":S0},
            "summary":summary, "adaptive_hard_gamma":{"median":float(np.median(gammas)),
            "mean":float(np.mean(gammas)),"min":float(np.min(gammas)),"max":float(np.max(gammas)),
            "median_score":float(np.median(scores)),"median_H":float(np.median(hvals))}}


if __name__ == "__main__":
    out = run()
    print(json.dumps(out, indent=2))
    Path("results.json").write_text(json.dumps(out, indent=2))