import json, math, random
from pathlib import Path
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score

SEED = 1021
rng = np.random.default_rng(SEED)
random.seed(SEED)


def certificate_sweep():
    # A recurrence whose exact tail has the paper's kappa^(l/3) form.
    eps = 1e-3
    C = 1.0
    rows = []
    for kappa in [0.05, 0.15, 0.30, 0.50, 0.70, 0.85, 0.95]:
        rho = kappa ** (1.0 / 3.0)
        # E_l is the remaining geometric effect after depth l.
        exact = lambda l: rho ** (l + 1) / (1.0 - rho)
        observed_ratio = exact(20) / exact(19)
        predicted_ratio = rho
        observed_l = next(l for l in range(10000) if exact(l) <= eps)
        # Formula in the proposal omits the geometric prefactor; use it literally.
        predicted_l = math.ceil(3 * math.log(C / eps) / math.log(1.0 / kappa) - 1)
        rows.append({"kappa": kappa, "predicted_tail_ratio": predicted_ratio,
                     "observed_tail_ratio": observed_ratio,
                     "predicted_depth": predicted_l, "observed_depth": observed_l})
    return rows


def boundary_sweep():
    # x_l=(gamma*Lambda)^l: predicted transition at gamma*Lambda=1.
    rows = []
    for q in [0.70, 0.90, 0.99, 1.00, 1.01, 1.10, 1.30]:
        x = q ** np.arange(30, dtype=float)
        # Estimate asymptotic multiplier from consecutive magnitudes.
        estimate = float(np.median(x[1:] / x[:-1]))
        rows.append({"gamma_times_Lambda": q, "observed_multiplier": estimate,
                     "predicted_regime": "decay" if q < 1 else ("critical" if q == 1 else "growth")})
    return rows


def make_sbm(n=600, p_in=0.045, p_out=0.006):
    y = np.repeat([0, 1], n // 2)
    A = np.zeros((n, n), dtype=np.float32)
    for i in range(n):
        same = y == y[i]
        prob = np.where(same, p_in, p_out)
        draw = rng.random(n) < prob
        draw[i] = False
        A[i, draw] = 1.0
    A = np.maximum(A, A.T)
    deg = A.sum(1)
    # symmetric normalized adjacency with self loops
    B = A + np.eye(n, dtype=np.float32)
    d = B.sum(1)
    P = B / np.sqrt(d[:, None] * d[None, :])
    X = np.zeros((n, 4), dtype=np.float32)
    X[:, 0] = (y * 2 - 1) + rng.normal(0, 1.0, n)
    X[:, 1] = (1 - 2 * y) + rng.normal(0, 1.0, n)
    X[:, 2:] = rng.normal(0, 1.0, (n, 2))
    return P.astype(np.float32), X, y


def gnn_experiment():
    P, X, y = make_sbm()
    n = len(y)
    perm = rng.permutation(n)
    train, val, test = perm[:240], perm[240:420], perm[420:]
    # Gamma is an attenuation estimate; Delta is the normalized propagation bound.
    gamma, Delta, eps, tau, max_depth = 0.72, 1.0, 0.02, 0.012, 8
    kappa = gamma * gamma * Delta
    # Calibrate C from a shallow-vs-deep validation error gap, as proposed.
    H = [X]
    for _ in range(max_depth):
        H.append(P @ H[-1])
    models = []
    val_errors = []
    for l in range(max_depth + 1):
        clf = LogisticRegression(C=2.0, max_iter=300, random_state=SEED)
        clf.fit(H[l][train], y[train])
        models.append(clf)
        val_errors.append(1 - accuracy_score(y[val], clf.predict(H[l][val])))
    gap = abs(val_errors[2] - val_errors[4])
    C = max(gap, 0.05)
    cert_depth = max(0, math.ceil(3 * math.log(C / eps) / math.log(1 / kappa) - 1))
    fixed_depth = min(max_depth, max(2, cert_depth))
    fixed_pred = models[fixed_depth].predict(H[fixed_depth][test])
    # Node-wise halting: two consecutive small logit changes plus certificate.
    logits = [m.predict_proba(H[l])[:, 1] for l, m in enumerate(models)]
    stop = np.full(n, max_depth, dtype=int)
    active = np.ones(n, dtype=bool)
    small_run = np.zeros(n, dtype=int)
    active_counts = []
    for l in range(1, max_depth + 1):
        d = np.abs(logits[l] - logits[l - 1])
        small_run = np.where(d < tau, small_run + 1, 0)
        tail = C * kappa ** ((l + 1) / 3)
        halt = active & (small_run >= 2) & (tail < eps)
        stop[halt] = l
        active[halt] = False
        active_counts.append(int(active.sum()))
    adaptive_prob = np.array([logits[stop[i]][i] for i in range(n)])
    adaptive_pred = (adaptive_prob >= 0.5).astype(int)
    # Full batched propagation does not realize node savings itself; this is the
    # equivalent active-root message count, useful for measuring the policy.
    fixed_work = len(test) * fixed_depth
    adaptive_work = int(stop[test].sum())
    return {
        "kappa": kappa, "calibrated_C": C, "certificate_depth": cert_depth,
        "fixed_depth": fixed_depth, "fixed_test_accuracy": float(accuracy_score(y[test], fixed_pred)),
        "adaptive_test_accuracy": float(accuracy_score(y[test], adaptive_pred[test])),
        "adaptive_mean_depth": float(stop[test].mean()),
        "work_ratio_adaptive_over_fixed": adaptive_work / fixed_work,
        "active_nodes_after_each_layer": active_counts,
        "validation_errors_by_depth": val_errors,
    }


def main():
    out = {"seed": SEED, "certificate_sweep": certificate_sweep(),
           "boundary_sweep": boundary_sweep(), "gnn": gnn_experiment()}
    Path("ks_results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()