import json
import numpy as np
from pathlib import Path

SEED = 3036
rng = np.random.default_rng(SEED)
N = 12
K = 1.0
DT = 0.08
STEPS = 250
QMAX = 3


def ring_adjacency(w):
    """Undirected ring with coupling w[d-1] at cyclic distance d=1,2."""
    A = np.zeros((N, N))
    for i in range(N):
        for d, wd in enumerate(w, 1):
            for j in ((i + d) % N, (i - d) % N):
                A[i, j] = wd
    return A


def prototype(q):
    return 2 * np.pi * q * np.arange(N) / N


def spectral_rate(A, alpha, psi):
    C = A * np.cos(psi[None, :] - psi[:, None] - alpha)
    L = np.diag(C.sum(axis=1)) - C
    ev = np.linalg.eigvals(L)
    # For a connected symmetric graph there is one gauge eigenvalue at zero.
    non_gauge = np.sort(np.real(ev[np.abs(ev) > 1e-8]))
    return float(K * non_gauge[0]), L, ev


def integrate_batch(theta, A, alpha):
    theta = theta.copy()
    for _ in range(STEPS):
        diff = theta[:, None, :] - theta[:, :, None]  # j-i, indexed [batch,i,j] after transpose below
        # Direct broadcasting: theta[b,j]-theta[b,i]-alpha[i,j]
        diff = theta[:, None, :] - theta[:, :, None] - alpha[None, :, :]
        theta += DT * K * np.sum(A[None, :, :] * np.sin(diff), axis=2)
        theta = (theta + np.pi) % (2 * np.pi) - np.pi
    return theta


def classify(theta):
    """Gauge-invariant cosine similarity to q-state prototypes."""
    z = np.exp(1j * theta)
    scores = []
    for q in range(QMAX + 1):
        p = np.exp(1j * prototype(q))
        # maximize over global phase by absolute complex overlap
        scores.append(np.abs(np.sum(z * np.conj(p), axis=1)) / N)
    scores = np.stack(scores, axis=1)
    return np.argmax(scores, axis=1), np.max(scores, axis=1)


def basin_fractions(w, samples=3000):
    A = ring_adjacency(w)
    alpha = np.zeros((N, N))
    theta0 = rng.uniform(-np.pi, np.pi, size=(samples, N))
    final = integrate_batch(theta0, A, alpha)
    labels, confidence = classify(final)
    counts = np.bincount(labels, minlength=QMAX + 1)
    return (counts / samples).tolist(), float(np.mean(confidence)), labels


def math_check():
    w = np.array([1.0, 0.55])
    A = ring_adjacency(w)
    alpha = np.zeros((N, N))
    psi = prototype(1)
    r, L, ev = spectral_rate(A, alpha, psi)
    # Jacobian of phase perturbations is -K L; measure decay in a non-gauge eigenmode.
    vals, vecs = np.linalg.eigh(L)
    ix = np.where(vals > 1e-8)[0][0]
    u = vecs[:, ix]
    eps = 1e-4
    theta = psi + eps * u
    decay = []
    for _ in range(80):
        f = K * np.sum(A * np.sin(theta[None, :] - theta[:, None]), axis=1)
        theta = theta + 0.01 * f
        err = theta - psi
        err -= err.mean()
        decay.append(np.linalg.norm(err))
    fitted = np.polyfit(np.arange(20, 75) * 0.01, np.log(np.maximum(decay[20:75], 1e-30)), 1)[0]
    return {"r_formula": r, "r_jacobian": float(K * vals[ix]),
            "measured_log_slope": float(fitted),
            "relative_rate_error": float(abs(r - K * vals[ix]) / max(abs(r), 1e-12))}


def main():
    check = math_check()
    # The q=1 phase pattern benefits from short-range coupling because its
    # edge phase difference is smaller there. Keep total coupling fixed.
    baseline = np.array([0.50, 0.50])
    shaped = np.array([0.90, 0.10])
    alpha = np.zeros((N, N))
    rates = {}
    for name, w in [("baseline", baseline), ("spectral_shaped", shaped)]:
        A = ring_adjacency(w)
        rates[name] = [spectral_rate(A, alpha, prototype(q))[0] for q in range(QMAX + 1)]
    b0, c0, _ = basin_fractions(baseline)
    b1, c1, _ = basin_fractions(shaped)
    # Across independent graphs, test the claimed positive rate/basin relation.
    corrs = []
    graph_rows = []
    for g in range(20):
        # random positive distance weights, normalized to the same total
        w = rng.dirichlet([2.0, 2.0])
        A = ring_adjacency(w)
        rr = np.array([spectral_rate(A, alpha, prototype(q))[0] for q in range(QMAX + 1)])
        bf, conf, _ = basin_fractions(w, samples=700)
        valid = np.isfinite(rr) & (np.array(bf) > 0)
        if valid.sum() > 1:
            corrs.append(float(np.corrcoef(rr[valid], np.array(bf)[valid])[0, 1]))
        graph_rows.append({"weights": w.tolist(), "rates": rr.tolist(), "basins": bf})
    result = {"seed": SEED, "N": N, "steps": STEPS, "dt": DT,
              "math_check": check,
              "baseline": {"weights": baseline.tolist(), "rates": rates["baseline"], "basins": b0, "mean_similarity": c0},
              "idea": {"weights": shaped.tolist(), "rates": rates["spectral_shaped"], "basins": b1, "mean_similarity": c1},
              "target_q": 1, "cross_graph_rate_basin_correlations": corrs,
              "cross_graph_mean_correlation": float(np.mean(corrs)),
              "cross_graph_median_correlation": float(np.median(corrs)),
              "graphs": graph_rows}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps({k: result[k] for k in ["math_check", "baseline", "idea", "cross_graph_mean_correlation", "cross_graph_median_correlation"]}, indent=2))

if __name__ == "__main__":
    main()