import json
import numpy as np

SEED = 1053
rng = np.random.default_rng(SEED)

# Two positive-definite phase Hessians. They do not commute.
H1 = np.array([[4.0, 1.0], [1.0, 1.0]])
H2 = np.array([[1.5, -0.7], [-0.7, 3.0]])
I = np.eye(2)


def monodromy(scale, order=(1, 2)):
    # eta1=1.6s and eta2=.4s: same two-step average LR s as baseline.
    e1, e2 = 1.6 * scale, 0.4 * scale
    J1, J2 = I - e1 * H1, I - e2 * H2
    return (J2 @ J1) if order == (1, 2) else (J1 @ J2)


def rho(M):
    return float(np.max(np.abs(np.linalg.eigvals(M))))


def boundary():
    # Largest stable scale, found from the exact rho(M)=1 crossing.
    lo, hi = 0.0, 2.0
    while rho(monodromy(hi)) < 1.0:
        hi *= 2
    for _ in range(70):
        mid = (lo + hi) / 2
        if rho(monodromy(mid)) < 1:
            lo = mid
        else:
            hi = mid
    return (lo + hi) / 2


def growth_check(s, periods=80):
    M = monodromy(s)
    vals, vecs = np.linalg.eig(M)
    j = int(np.argmax(np.abs(vals)))
    v = np.real(vecs[:, j])
    v /= np.linalg.norm(v)
    x = v.copy()
    norms = []
    for _ in range(periods):
        x = M @ x
        norms.append(np.linalg.norm(x))
    # Geometric mean removes initial normalization and reports empirical Floquet growth.
    empirical = float((norms[-1] / norms[0]) ** (1.0 / (periods - 1)))
    return rho(M), empirical, float(norms[-1])


def commutator_scaling():
    C = H2 @ H1 - H1 @ H2
    c_norm = np.linalg.norm(C, 2)
    rows = []
    for s in [0.01, 0.02, 0.04, 0.08, 0.16]:
        # Exact antisymmetric part is eta1*eta2*(H2 H1-H1 H2).
        M = monodromy(s)
        observed = np.linalg.norm(M - M.T, 2)
        predicted = (1.6 * s) * (0.4 * s) * c_norm
        rows.append({"scale": s, "observed": observed, "predicted": predicted,
                     "ratio": observed / predicted})
    return {"commutator_norm": c_norm, "rows": rows}


def quadratic_optimizer(scale=0.20, periods=60):
    # Both methods see the same alternating phase gradients. Baseline uses constant
    # LR s; Floquet uses the two phase LRs while preserving the same average LR.
    x0 = np.array([2.0, -1.5])
    x_const, x_periodic = x0.copy(), x0.copy()
    e1, e2 = 1.6 * scale, .4 * scale
    trace = []
    Hbar = .5 * (H1 + H2)
    def objective(x):
        return float(.5 * x @ Hbar @ x)
    for p in range(periods):
        # phase 1
        x_const -= scale * (H1 @ x_const)
        x_periodic -= e1 * (H1 @ x_periodic)
        # phase 2
        x_const -= scale * (H2 @ x_const)
        x_periodic -= e2 * (H2 @ x_periodic)
        if p in (0, 4, 9, 19, 39, periods - 1):
            trace.append({"period": p + 1, "constant_loss": objective(x_const),
                          "periodic_loss": objective(x_periodic)})
    return {"scale": scale, "periods": periods, "trace": trace,
            "final_constant_loss": objective(x_const),
            "final_periodic_loss": objective(x_periodic)}


def main():
    s_star = boundary()
    # Boundary prediction is the formula rho[(I-e2 H2)(I-e1 H1)]=1;
    # observed transition is independently classified by long-run norm behavior.
    scales = [0.90 * s_star, 0.99 * s_star, 1.01 * s_star, 1.10 * s_star]
    growth = []
    for s in scales:
        r, empirical, final_norm = growth_check(s)
        growth.append({"scale": s, "rho_predicted": r,
                       "empirical_per_period_growth": empirical,
                       "final_perturbation_norm": final_norm,
                       "classification": "stable" if empirical < 1 else "unstable"})
    # A direct sweep gives an independently observable transition bracket.
    sweep = []
    for s in np.linspace(.8 * s_star, 1.2 * s_star, 17):
        r, empirical, _ = growth_check(float(s), periods=50)
        sweep.append([float(s), r, empirical])
    order_difference = np.linalg.norm(monodromy(.25, (1, 2)) - monodromy(.25, (2, 1)))
    result = {
        "seed": SEED,
        "hessians": {"H1": H1.tolist(), "H2": H2.tolist()},
        "predictions": {
            "critical_scale_exact": s_star,
            "stability_rule": "rho(M)<1 predicts decay and rho(M)>1 predicts growth",
            "commutator_rule": "||M-M.T|| = eta1*eta2*||[H2,H1]|| for symmetric H phases"
        },
        "growth_near_boundary": growth,
        "boundary_sweep": sweep,
        "commutator_scaling": commutator_scaling(),
        "reversed_order_matrix_difference_at_scale_.25": order_difference,
        "quadratic_optimizer": quadratic_optimizer()
    }
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()