import json
import numpy as np
from pathlib import Path


def stats(k, p):
    k = np.asarray(k, float); p = np.asarray(p, float)
    mean_k = p @ k
    return mean_k, p @ (k*k), (p @ (k*k))/mean_k


def order_parameter(m, p):
    mbar = p @ m
    return float(p @ ((m - mbar)**2)), float(mbar)


def jacobian(m, k, p, gain, reg=0.0, mode="degree"):
    # gain = beta*J; this is the exact Jacobian of the supplied ODE.
    mean_k = p @ k
    z = gain * k * ((p * k) @ m) / mean_k
    d = 1.0 - np.tanh(z)**2
    A = -np.eye(len(k)) + gain * d[:, None] * (k[:, None] * (p*k)[None, :] / mean_k)
    if reg:
        if mode == "degree":
            A -= 2.0 * reg * (np.eye(len(k)) - np.ones((len(k), 1)) @ p[None, :])
        elif mode == "global":
            A -= 2.0 * reg * np.eye(len(k))
    return A


def integrate(k, p, gain, reg=0.0, mode="none", seed=0, steps=6000, dt=.05):
    rng = np.random.default_rng(seed)
    # Small degree-correlated perturbation selects one of the pitchfork branches.
    m = 1e-3 * (k - p @ k) + 1e-5 * rng.normal(size=len(k))
    mean_k = p @ k
    for _ in range(steps):
        u = (p*k) @ m / mean_k
        dm = -m + np.tanh(gain*k*u)
        if reg == 0 or mode == "none":
            pass
        elif mode == "degree":
            dm -= 2.0*reg*(m - p @ m)
        elif mode == "global":
            dm -= 2.0*reg*m
        m += dt*dm
        if not np.all(np.isfinite(m)):
            raise FloatingPointError("nonfinite state")
    S, mb = order_parameter(m, p)
    return m, S, mb


def restricted_eigenvalue(A, p):
    # Diagnostic specified in the idea: project onto p-weighted zero-mean vectors.
    n = len(p)
    P = np.eye(n) - np.ones((n, 1)) @ p[None, :]
    # Nullspace basis of p^T using SVD; projected operator is represented there.
    _, _, vh = np.linalg.svd(p[None, :])
    Q = vh[1:].T
    vals = np.linalg.eigvals(Q.T @ P @ A @ Q)
    return float(np.max(np.real(vals)))


def run():
    rows = []
    # Equal mean degree, increasing heterogeneity: p=(1/2,1/2), k=(3-d,3+d).
    gains = np.linspace(.08, .70, 32)
    for d in [0.0, 1.0, 2.0, 4.0]:
        k = np.array([3.0-d, 3.0+d]); p = np.array([.5, .5])
        mk, mk2, ratio = stats(k,p)
        predicted = 1.0/ratio
        for g in gains:
            m,S,mb = integrate(k,p,g,seed=123)
            A = jacobian(np.zeros(2),k,p,g)
            rows.append(dict(kind="sweep", heterogeneity=d, gain=float(g), S=S,
                             mean=float(mb), lambda_full=float(np.max(np.linalg.eigvals(A).real)),
                             lambda_sep=restricted_eigenvalue(A,p), predicted_gain=predicted,
                             k2_over_k=ratio))
    # Regularizer comparison at a supercritical gain.
    k=np.array([1.,5.]); p=np.array([.5,.5]); g=.55
    for mode, reg in [("none",0.),("global",.20),("degree",.20),("degree",.50)]:
        m,S,mb=integrate(k,p,g,reg=reg,mode=mode,seed=123)
        rows.append(dict(kind="regularizer", mode=mode, reg=reg, gain=g, S=S, mean=mb,
                         norm=float(np.sqrt(p@(m*m)))))
    out=Path("results.json")
    out.write_text(json.dumps(rows, indent=2))
    # Compact summary used by the report and easy reproduction.
    summary=[]
    for d in [0.,1.,2.,4.]:
        sub=[r for r in rows if r.get("kind")=="sweep" and r["heterogeneity"]==d]
        pred=sub[0]["predicted_gain"]
        # first sampled point whose S exceeds a robust numerical threshold
        active=[r["gain"] for r in sub if r["S"]>1e-8]
        summary.append({"heterogeneity":d,"predicted":pred,"observed_grid_onset":min(active) if active else None,
                        "max_S":max(r["S"] for r in sub)})
    print(json.dumps({"threshold_summary":summary,
                      "regularizers":[r for r in rows if r.get("kind")=="regularizer"]}, indent=2))


if __name__ == "__main__":
    run()