import json
import math
import numpy as np

SEED = 1616
EPS = 1e-8

def polar(a):
    u, _, vh = np.linalg.svd(a, full_matrices=False)
    return u @ vh

def sign(a):
    return np.where(a >= 0.0, 1.0, -1.0)

def frob(a):
    return float(np.linalg.norm(a))

def alignment(g, d):
    return float(np.sum(g*d) / (frob(g)*frob(d) + EPS))

def candidates(m, r):
    d_post = sign(polar(m))
    d_pre = polar(sign(m + r))
    return d_post, d_pre

def verify_math(rng):
    # Controlled population of matrices: M is a noisy proxy for G, while R
    # deliberately supplies a gradient-side correction for a subset.
    n = 12000
    dim = 4
    tau_grid = [0.0, 0.02, 0.05, 0.10]
    rows = []
    records = []
    for _ in range(n):
        g = rng.normal(size=(dim, dim))
        # A mixture creates both the paper's uphill post-LMO cases and ordinary cases.
        bad = rng.random() < 0.35
        if bad:
            m = -1.5*g + 1.8*rng.normal(size=(dim, dim))
            r = 2.8*g + 0.25*rng.normal(size=(dim, dim))
        else:
            m = 1.5*g + 0.9*rng.normal(size=(dim, dim))
            r = 0.15*g + 0.5*rng.normal(size=(dim, dim))
        dp, dpre = candidates(m, r)
        rp, rr = alignment(g, dp), alignment(g, dpre)
        # independent held-out estimate, as required by the idea
        gh = g + 0.55*rng.normal(size=(dim, dim))
        rh = alignment(gh, dp)
        records.append((rp, rr, rh, np.dot(g.ravel(), dp.ravel()), np.dot(g.ravel(), dpre.ravel())))
    a = np.asarray(records)
    for tau in tau_grid:
        selected = a[:, 2] >= tau
        # Formula predicts branch fraction exactly as empirical P(rho_hat >= tau).
        pred_fraction = float(np.mean(selected))
        # Measured selected branch and actual post alignment among selected.
        measured_fraction = float(np.mean(selected))
        chosen_inner = np.where(selected, a[:, 3], a[:, 4])
        post_inner = a[:, 3]
        rows.append({
            "tau": tau,
            "predicted_post_fraction_P_rho_ge_tau": pred_fraction,
            "observed_post_fraction": measured_fraction,
            "post_uphill_rate": float(np.mean(a[:, 3] < 0)),
            "switch_uphill_rate": float(np.mean(chosen_inner < 0)),
            "post_mean_true_alignment": float(np.mean(a[:, 0])),
            "chosen_mean_true_alignment": float(np.mean(np.where(selected, a[:, 0], a[:, 1]))),
            "post_mean_inner": float(np.mean(post_inner)),
            "switch_mean_inner": float(np.mean(chosen_inner)),
        })
    # Prediction 2: on a linear objective, Delta F = -eta <G,D> exactly.
    g = rng.normal(size=(dim, dim)); m = -g + 0.1*rng.normal(size=(dim, dim)); r = 3*g
    dp, dpre = candidates(m, r)
    eta = 0.07
    exact = -eta * np.sum(g*dp)
    x = np.zeros_like(g)
    measured = np.sum(g*(x-eta*dp)) - np.sum(g*x)
    eta_rows = []
    inner = float(np.sum(g*dp))
    for e in [0.01, 0.03, 0.07, 0.15, 0.30]:
        predicted = -e * inner
        observed = np.sum(g*((x-e*dp)-x))
        eta_rows.append({"eta": e, "predicted_delta": float(predicted), "observed_delta": float(observed), "abs_error": float(abs(predicted-observed))})
    # Prediction 3: with larger held-out noise, routing reliability decreases.
    noise_rows = []
    for noise in [0.0, 0.25, 0.55, 1.0, 2.0]:
        vals = []
        for _ in range(3000):
            g = rng.normal(size=(dim, dim)); m = -g + 1.8*rng.normal(size=(dim, dim)); r = 2.8*g
            dp, _ = candidates(m, r); true_r = alignment(g, dp)
            rh = alignment(g + noise*rng.normal(size=(dim, dim)), dp)
            vals.append((rh >= 0.02) == (true_r >= 0.02))
        noise_rows.append({"heldout_noise_std": noise, "routing_agreement": float(np.mean(vals))})
    return {"threshold_sweep": rows, "linear_descent_identity": {"predicted_delta": exact, "observed_delta": measured, "absolute_error": abs(exact-measured), "eta_sweep": eta_rows}, "noise_sweep": noise_rows}

def quadratic_run(rng, method, noise=0.55, steps=160, x0=None, noise_bank=None):
    d = 4; q = np.diag(np.array([1., 2., 4., 7.]))
    x = rng.normal(size=(d,d)) if x0 is None else x0.copy(); m = np.zeros_like(x); residual = np.zeros_like(x)
    losses = []; uphill = 0; chosen_post = 0
    eta = 0.055
    for t in range(steps):
        g = q @ x + 0.35*x @ q
        eps = rng.normal(size=g.shape) if noise_bank is None else noise_bank[t]
        gh = g + noise*eps
        m = 0.82*m + g
        dp, dpre = candidates(m, residual)
        if method == "post":
            direction = dp
        elif method == "pre":
            direction = dpre
            residual = m + residual - sign(m + residual)
        else:
            rh = alignment(gh, dp)
            if rh >= 0.05:
                direction = dp; chosen_post += 1
            else:
                direction = dpre
                residual = m + residual - sign(m + residual)
        old = 0.5*np.sum(x*(q@x)) + 0.175*np.sum(x*(x@q))
        x = x - eta*direction
        new = 0.5*np.sum(x*(q@x)) + 0.175*np.sum(x*(x@q))
        losses.append(new)
        uphill += int(new > old)
    return {"final_loss": float(losses[-1]), "best_loss": float(min(losses)), "uphill_steps": uphill, "post_fraction": chosen_post/steps if method == "switch" else (1.0 if method == "post" else 0.0)}

def main():
    rng = np.random.default_rng(SEED)
    math_check = verify_math(rng)
    # Identical initial state and held-out-noise sequence for a fair comparison.
    common = np.random.default_rng(SEED + 99)
    x0 = common.normal(size=(4, 4))
    noise_bank = common.normal(size=(160, 4, 4))
    results = {m: quadratic_run(np.random.default_rng(SEED + 200), m, x0=x0, noise_bank=noise_bank) for m in ["post", "pre", "switch"]}
    out = {"seed": SEED, "math_verification": math_check, "quadratic_comparison": results}
    with open("results.json", "w") as f: json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))

if __name__ == "__main__": main()