import json
from pathlib import Path
import numpy as np

# Harmonic-mode branch: minimal numerical verification of the stated decomposition.
# The synthetic PSD operator is equivalent to a Hodge Laplacian in an orthonormal basis:
# L = Q diag(0,...,0, positive eigenvalues) Q^T.

def make_operator(seed=7, dim=12, harmonic_dim=2):
    rng = np.random.default_rng(seed)
    q, _ = np.linalg.qr(rng.normal(size=(dim, dim)))
    positive = np.array([0.35, 0.8, 1.7, 2.6, 3.9, 5.2, 6.4, 7.5, 9.0, 11.0])[: dim-harmonic_dim]
    eig = np.concatenate([np.zeros(harmonic_dim), positive])
    L = q @ np.diag(eig) @ q.T
    H = q[:, :harmonic_dim]
    P_h = H @ H.T
    P_p = np.eye(dim) - P_h
    return L, H, P_h, P_p, eig, q


def fit_log_slope(values, dt):
    # Ignore the t=0 point and fit where values are safely above roundoff.
    vals = np.asarray(values)
    t = np.arange(len(vals)) * dt
    keep = (vals > 1e-11) & (np.arange(len(vals)) > 2)
    return float(np.polyfit(t[keep], np.log(vals[keep]), 1)[0])


def run():
    L, H, P_h, P_p, eig, Q = make_operator()
    dim, k = L.shape[0], H.shape[1]
    lam_pos = float(np.min(eig[eig > 0]))
    lam_max = float(np.max(eig))
    rng = np.random.default_rng(123)
    h0 = H @ np.array([1.25, -0.7])
    # Exact slowest dissipative mode makes the continuous-rate prediction sharp.
    u0 = Q[:, k] * 1.0
    z0 = h0 + u0

    # Prediction 1: projected branch has constant harmonic coordinates, while a
    # fully damped baseline incorrectly erases them.
    nu = 0.9
    dt = 0.01
    steps = 900
    z_split = z0.copy()
    z_full = z0.copy()
    harmonic_errors = []
    for _ in range(steps):
        # N is identically zero here; projected nonlinear N would preserve the same invariant.
        z_split = z_split + dt * (-nu * (L @ (P_p @ z_split)))
        z_split = P_h @ z_split + P_p @ z_split  # explicit numerical projection
        z_full = z_full + dt * (-nu * (L @ z_full) - nu * (P_h @ z_full))
        harmonic_errors.append(np.linalg.norm(P_h @ (z_split - z0)))
    invariant_error = float(max(harmonic_errors))
    baseline_memory_retention = float(np.linalg.norm(P_h @ z_full) / np.linalg.norm(h0))
    split_memory_retention = float(np.linalg.norm(P_h @ z_split) / np.linalg.norm(h0))

    # Prediction 2: dissipative amplitude has log slope -nu*lambda_1^+.
    nus = [0.25, 0.5, 0.9, 1.4, 2.0]
    slope_rows = []
    for n in nus:
        x = u0.copy()
        amps = []
        for _ in range(500):
            amps.append(np.linalg.norm(x))
            x = x + dt * (-n * (L @ x))
        observed = fit_log_slope(amps, dt)
        predicted = -n * lam_pos
        slope_rows.append({"nu": n, "predicted": predicted, "observed": observed,
                           "abs_error": abs(observed - predicted)})

    # Prediction 3: explicit Euler is stable iff dt*nu*lambda_max < 2 for PSD L.
    n_stab = 1.0
    predicted_dt = 2.0 / (n_stab * lam_max)
    factors = [0.90, 0.99, 1.01, 1.10]
    stability_rows = []
    x_init = Q[:, -1]
    for f in factors:
        d = f * predicted_dt
        x = x_init.copy()
        norms = []
        for _ in range(80):
            norms.append(float(np.linalg.norm(x)))
            x = x - d * n_stab * (L @ x)
        # Stable side remains bounded and decays; unstable side grows geometrically.
        observed_stable = max(norms) < 10.0 and norms[-1] <= norms[0]
        stability_rows.append({"factor_of_boundary": f, "dt": d,
                               "predicted": "stable" if f < 1 else "unstable",
                               "observed": "stable" if observed_stable else "unstable",
                               "final_norm": norms[-1]})

    # Small trajectory comparison: standard/full damping vs topology-preserving split.
    # Both use exactly the same operator, timestep, and initialization.
    comparison = {
        "split_harmonic_retention": split_memory_retention,
        "fully_damped_harmonic_retention": baseline_memory_retention,
        "split_dissipative_final_norm": float(np.linalg.norm(P_p @ z_split)),
        "fully_damped_final_norm": float(np.linalg.norm(z_full)),
    }
    result = {
        "operator_dim": dim, "harmonic_dim": k, "lambda_1_positive": lam_pos,
        "lambda_max": lam_max, "predicted_euler_boundary_dt": predicted_dt,
        "prediction_1_invariant_max_error": invariant_error,
        "prediction_1_tolerance": 1e-12,
        "prediction_2_decay_sweep": slope_rows,
        "prediction_3_stability_sweep": stability_rows,
        "comparison": comparison,
        "claims": {
            "harmonic_invariant": invariant_error < 1e-12 and split_memory_retention > 0.999999,
            "decay_scaling": max(r["abs_error"] for r in slope_rows) < 0.01,
            "euler_boundary": all(r["observed"] == r["predicted"] for r in stability_rows),
        }
    }
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    run()