import json
from pathlib import Path
import numpy as np


def schur_monitor(H, G, lam=1e-3, beta=1.0):
    """Return R=G.T (H+lam I)^-1 G and the trust update operator."""
    H = np.asarray(H, dtype=float)
    G = np.asarray(G, dtype=float)
    A = H + lam * np.eye(H.shape[0])
    X = np.linalg.solve(A, G)
    R = 0.5 * (G.T @ X + X.T @ G)
    # eigh is stable and makes tiny numerical asymmetry harmless
    ev, V = np.linalg.eigh(R)
    ev = np.maximum(ev, 0.0)
    R_psd = (V * ev) @ V.T
    trust = np.linalg.inv(np.eye(G.shape[1]) + beta * R_psd)
    return R_psd, trust


def reduced_energy(sigma, u, H, G, a):
    # E is affine in u before relaxation: .5 sigma'Hsigma + sigma'Gu + a'u
    return .5 * sigma @ H @ sigma + sigma @ G @ u + a @ u


def verify_identity(rng):
    k, m = 4, 3
    A = rng.normal(size=(k, k))
    H = A.T @ A + 0.7 * np.eye(k)
    G = rng.normal(size=(k, m))
    a = rng.normal(size=m)
    HinvG = np.linalg.solve(H, G)
    R = G.T @ HinvG
    # Relax sigma*(u)=-H^-1 G u and finite-difference the reduced energy.
    def reduced(u):
        sigma = -np.linalg.solve(H, G @ u)
        return reduced_energy(sigma, u, H, G, a)
    u0 = rng.normal(size=m)
    eps = 2e-4
    numerical = np.zeros((m, m))
    for i in range(m):
        ei = np.eye(m)[i]
        for j in range(m):
            ej = np.eye(m)[j]
            numerical[i, j] = (reduced(u0+eps*ei+eps*ej)-reduced(u0+eps*ei-eps*ej)
                - reduced(u0-eps*ei+eps*ej)+reduced(u0-eps*ei-eps*ej))/(4*eps**2)
    target = -R
    eig_R = np.linalg.eigvalsh(R)
    return {
        "max_abs_hessian_error": float(np.max(np.abs(numerical-target))),
        "R_eigenvalues": eig_R.tolist(),
        "min_R_eigenvalue": float(eig_R.min()),
        "identity_pass": bool(np.max(np.abs(numerical-target)) < 2e-6 and eig_R.min() >= -1e-10),
    }


def run_control(rng, steps=80):
    # Two mechanism amplitudes. R has one very interactive direction.
    H = np.diag([0.25, 1.0])
    G = np.array([[2.7, 2.7], [0.15, -0.15]])
    R, trust = schur_monitor(H, G, lam=0.02, beta=1.0)
    # Smooth bounded target loss, with a deliberately aggressive hyper-step.
    target = np.array([1.0, -1.0])
    Q = np.diag([1.0, 1.0])
    eta = 0.72
    rows = {"baseline": [], "schur": []}
    noise = 0.015 * rng.normal(size=(steps, 2))
    for name, P in [("baseline", np.eye(2)), ("schur", trust)]:
        u = np.array([0.0, 0.0])
        for t in range(steps):
            # Hypergradient of a quadratic validation proxy, plus reproducible mild noise.
            h = Q @ (u-target) + noise[t]
            delta = -eta * P @ h
            u = np.clip(u + delta, -3.0, 3.0)
            loss = 0.5 * (u-target) @ Q @ (u-target)
            rows[name].append({"step": t+1, "loss": float(loss), "u_norm": float(np.linalg.norm(u)),
                               "delta_norm": float(np.linalg.norm(delta))})
    def summarize(x):
        losses = np.array([z["loss"] for z in x])
        return {"final_loss": float(losses[-1]), "best_loss": float(losses.min()),
                "loss_spikes": int(np.sum(losses[1:] > 1.2 * losses[:-1])),
                "mean_last10": float(losses[-10:].mean()),
                "max_u_norm": float(max(z["u_norm"] for z in x))}
    return {"H": H.tolist(), "G": G.tolist(), "R": R.tolist(),
            "R_eigenvalues": np.linalg.eigvalsh(R).tolist(),
            "trust_matrix": trust.tolist(), "eta": eta,
            "summary": {k: summarize(v) for k,v in rows.items()}, "trace": rows}


def main():
    rng = np.random.default_rng(491)
    verification = verify_identity(rng)
    experiment = run_control(rng)
    out = {"verification": verification, "experiment": experiment}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps({"verification": verification, "summary": experiment["summary"],
                      "R_eigenvalues": experiment["R_eigenvalues"]}, indent=2))

if __name__ == "__main__":
    main()