"""Directed-path synchronization coupling: math checks and a tiny experiment.

Convention: C has C[i,j]=a_ij for j != i and C[i,i]=-sum_j a_ij,
so xdot = f(x) - eta*C_pos?  We instead construct the positive
row Laplacian B with B[i,i]=sum a and B[i,j]=-a, and use
xdot = f(x) - eta B x.  Every non-root node receives its parent.
"""
import json
import numpy as np

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


def directed_star(n, weight=1.0):
    # n-1 directed parent paths (root 0 -> each child), represented as
    # receiver rows: child receives the root state.
    B = np.zeros((n, n))
    for i in range(1, n):
        B[i, i] = weight
        B[i, 0] = -weight
    return B


def dense_symmetric(n, weight=1.0):
    # Complete balanced diffusive graph, same row-sum normalization.
    B = np.full((n, n), -weight / (n - 1))
    np.fill_diagonal(B, weight)
    return B


def chain(n, weights):
    # Imbalanced directed path: each node receives only its predecessor.
    B = np.zeros((n, n))
    for i, w in enumerate(weights, start=1):
        B[i, i] = w
        B[i, i - 1] = -w
    return B


def graph_gain(B, tol=1e-9):
    """Conservative spectral gain: smallest positive real Laplacian mode."""
    ev = np.linalg.eigvals(B)
    positive = [z.real for z in ev if z.real > tol and abs(z.imag) < 1e-7]
    if not positive:
        # For general directed matrices, use the minimum positive real part.
        positive = [z.real for z in ev if z.real > tol]
    return float(min(positive))


def simulate(B, lam, eta, dt=0.002, T=5.0, x0=None):
    n = B.shape[0]
    steps = int(T / dt)
    x = np.array(x0 if x0 is not None else rng.normal(size=n), dtype=float)
    # For directed graphs use the left-null weighted consensus coordinate.
    left = np.linalg.svd(B.T)[2][-1]
    left = left / np.sum(left)
    V = []
    for _ in range(steps + 1):
        mean = np.sum(left * x)
        V.append(0.5 * np.sum((x - mean) ** 2))
        x = x + dt * (lam * x - eta * (B @ x))
    return np.asarray(V)


def slope(V, dt, tail_fraction=0.6):
    k = np.arange(len(V))
    first = int(len(V) * (1-tail_fraction))
    y = np.log(np.maximum(V, 1e-300))
    return float(np.polyfit(k[first:] * dt, y[first:], 1)[0])


def threshold_sweep(B, lam, gamma, etas):
    # Stable means final V is below its initial V, with a substantial margin.
    x0 = np.linspace(-1.0, 1.0, B.shape[0])
    rows = []
    for eta in etas:
        V = simulate(B, lam, eta, x0=x0)
        rows.append({"eta": float(eta), "final_over_initial": float(V[-1]/V[0]),
                     "slope": slope(V, .002)})
    return rows


def main():
    n = 8
    lam = 0.8                 # exact Lipschitz constant of f(x)=lam*x
    star = directed_star(n)
    dense = dense_symmetric(n)
    imbalanced = chain(n, [0.35, 0.55, 0.75, 0.95, 1.15, 1.35, 1.55])
    graphs = {"directed_n_minus_1_star": star, "dense_symmetric": dense,
              "severely_imbalanced_path": imbalanced}
    gains = {name: graph_gain(B) for name, B in graphs.items()}
    # Prediction 1: boundary eta_c = L/gamma.
    boundary = {name: lam/g for name, g in gains.items()}
    etas = np.array([0.3, 0.6, 0.79, 0.81, 1.0, 1.4, 2.0])
    sweeps = {name: threshold_sweep(B, lam, gains[name], etas)
              for name, B in graphs.items()}

    # Prediction 2: rate d log(V)/dt approximately -2(eta*gamma-L).
    rate_rows = []
    for eta in [1.0, 1.4, 2.0]:
        V = simulate(star, lam, eta, x0=np.linspace(-1, 1, n))
        observed = slope(V, .002)
        predicted = -2*(eta*gains["directed_n_minus_1_star"] - lam)
        rate_rows.append({"eta": eta, "predicted": predicted, "observed": observed,
                          "relative_error": abs(observed-predicted)/max(abs(predicted),1e-9)})

    # Prediction 3: threshold scales linearly with L (gamma fixed).
    scaling = []
    for L in [0.2, 0.5, 0.8, 1.2]:
        g = gains["directed_n_minus_1_star"]
        candidate = L/g
        # Estimate the threshold from the first negative fitted log-V slope.
        candidates = np.linspace(max(.02, candidate*.5), candidate*1.5, 61)
        found = None
        for eta in candidates:
            V = simulate(star, L, eta, x0=np.linspace(-1, 1, n))
            if slope(V, .002) < 0:
                found = float(eta); break
        scaling.append({"L": L, "predicted_eta_c": candidate,
                        "observed_negative_slope_eta": found,
                        "relative_boundary_error": (abs(found-candidate)/candidate
                            if found is not None else None)})

    # Standard baselines at the same lambda and eta.
    x0 = np.linspace(-1, 1, n)
    baseline = simulate(np.zeros((n,n)), lam, 1.0, x0=x0)
    idea = simulate(star, lam, 1.0, x0=x0)
    dense_run = simulate(dense, lam, 1.0, x0=x0)
    out = {"seed": SEED, "L": lam, "graph_gains": gains,
           "predicted_thresholds": boundary, "threshold_sweep": sweeps,
           "rate_test": rate_rows, "L_scaling": scaling,
           "comparison_eta_1": {
               "uncoupled_final_over_initial": float(baseline[-1]/baseline[0]),
               "directed_path_final_over_initial": float(idea[-1]/idea[0]),
               "dense_symmetric_final_over_initial": float(dense_run[-1]/dense_run[0]),
               "uncoupled_slope": slope(baseline,.002),
               "directed_path_slope": slope(idea,.002),
               "dense_symmetric_slope": slope(dense_run,.002)}}
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()