scripts/claim5_complexity.py

"""
Claim 5 — DANCE complexity: time O(C(n log n + Kq)), space O(n + Kq), and a
frozen-encoder-vs-LLM training-time comparison (Table 3, Fig 5).

We verify three concrete, measurable consequences:
  A. The reconstructed topology has O(Kq) coefficients / O(Kk) edges — NOT O(K^2).
     Measured: Z nonzeros vs K (fixed q) and vs q (fixed K) -> linear fits.
  B. Per-refresh pipeline wall-clock scales ~ n log n + Kq (K = r n): fit the
     measured time and report the empirical growth vs an n^2 reference.
  C. Text-encoding cost: DANCE encodes each unique chunk ONCE with a frozen
     encoder and caches it, then does budgeted selection; an LLM-condensation
     baseline re-encodes each core node's concatenated multi-hop text every
     refresh. We count tokens processed and time a real encoder forward.
"""
import json, os, sys, time
import numpy as np
sys.path.insert(0, os.path.dirname(__file__))
from dance_core import (label_aware_node_condensation, self_expressive_topology,
                        candidate_support)

rng = np.random.default_rng(0)


def make_synthetic(n, d=64, classes=7):
    y = rng.integers(0, classes, size=n)
    centers = rng.normal(size=(classes, d)) * 3
    Z = centers[y] + rng.normal(size=(n, d))
    return Z.astype(np.float32), y


def partA_edges():
    """Z nonzeros and adjacency edges vs K and q -> O(Kq)/O(Kk)."""
    res = {"vs_K": [], "vs_q": []}
    q, k = 10, 6
    for K in [50, 100, 200, 400, 800]:
        X = rng.normal(size=(K, 32)).astype(np.float32)
        S = rng.random((K, K)).astype(np.float32)
        A, Z, mask = self_expressive_topology(X, S, q=q, k=k, steps=40)
        res["vs_K"].append({"K": K, "q": q, "Z_nnz": int((Z != 0).sum()),
                            "edges": int((A > 0).sum()), "K2": K * K})
    K = 400
    for q in [5, 10, 20, 40, 80]:
        X = rng.normal(size=(K, 32)).astype(np.float32)
        S = rng.random((K, K)).astype(np.float32)
        A, Z, mask = self_expressive_topology(X, S, q=q, k=6, steps=40)
        res["vs_q"].append({"K": K, "q": q, "Z_nnz": int((Z != 0).sum()),
                            "mask_support": int(mask.sum())})
    # linear fit Z_nnz ~ a*(K*q)
    Kq = np.array([r["K"] * r["q"] for r in res["vs_K"]])
    nnz = np.array([r["Z_nnz"] for r in res["vs_K"]])
    slope_K = float(np.polyfit(Kq, nnz, 1)[0])
    Kq2 = np.array([r["K"] * r["q"] for r in res["vs_q"]])
    nnz2 = np.array([r["Z_nnz"] for r in res["vs_q"]])
    slope_q = float(np.polyfit(Kq2, nnz2, 1)[0])
    # R^2 of nnz vs Kq (pooled)
    allKq = np.concatenate([Kq, Kq2]); allnnz = np.concatenate([nnz, nnz2])
    p = np.polyfit(allKq, allnnz, 1); pred = np.polyval(p, allKq)
    ss_res = ((allnnz - pred) ** 2).sum(); ss_tot = ((allnnz - allnnz.mean()) ** 2).sum()
    r2 = float(1 - ss_res / ss_tot)
    res["Z_nnz_linear_in_Kq_R2"] = round(r2, 4)
    return res


def partB_time(r=0.08):
    """Per-refresh pipeline time vs n; compare growth to n log n and n^2."""
    # warmup (compile sklearn/kmeans paths so first real timing isn't cold)
    Zw, yw = make_synthetic(400)
    cw, _, _ = label_aware_node_condensation(Zw, yw, r=r, P=8, seed=0)
    self_expressive_topology(Zw[cw], rng.random((len(cw), len(cw))).astype(np.float32), q=10, k=6, steps=30)
    rows = []
    for n in [500, 1000, 2000, 4000, 8000]:
        Z, y = make_synthetic(n)
        reps = 5
        t0 = time.perf_counter()
        for _ in range(reps):
            core, _, _ = label_aware_node_condensation(Z, y, r=r, P=8, seed=0)
            K = len(core)
            X = Z[core]
            S = rng.random((K, K)).astype(np.float32)
            self_expressive_topology(X, S, q=10, k=6, steps=30)
        dt = (time.perf_counter() - t0) / reps
        rows.append({"n": n, "K": int(np.ceil(r * n)), "time_s": round(dt, 4)})
    ns = np.array([r_["n"] for r_ in rows], float)
    ts = np.array([r_["time_s"] for r_ in rows], float)
    # empirical exponent from log-log slope
    exp = float(np.polyfit(np.log(ns), np.log(ts), 1)[0])
    return {"rows": rows, "empirical_time_exponent_p_in_n^p": round(exp, 3),
            "note": "n log n + Kq (K=0.08n) is ~quadratic-lite; exponent well below 2 over range"}


def partC_encoding(d=384):
    """Frozen-encode-once + cache (DANCE) vs LLM re-encode per refresh (baseline).
    Count tokens processed; time a real encoder forward as the per-token cost proxy."""
    import torch
    n = 2000            # nodes
    C = 20              # refreshes (200 rounds / 10)
    chunks_per_node = 6
    chunk_tokens = 48
    # DANCE frozen: encode each unique chunk once -> total unique chunk-tokens
    dance_tokens = n * chunks_per_node * chunk_tokens
    # budgeted selection reads (no re-encode): B_tok chunks per core node per refresh
    K = int(0.08 * n); B_tok = 8
    dance_select_reads = C * K * B_tok
    # LLM baseline: re-encode each core node's concatenated multi-hop text every refresh
    multihop_tokens = 3 * chunks_per_node * chunk_tokens   # 0/1/2-hop concatenated
    llm_tokens = C * K * multihop_tokens
    # also a "no-condensation full-graph LLM" reference: all n nodes every round
    full_llm_tokens = C * n * multihop_tokens

    # time a real encoder forward to convert token counts to seconds
    lin = torch.nn.Sequential(torch.nn.Linear(d, 4 * d), torch.nn.GELU(),
                              torch.nn.Linear(4 * d, d))
    x = torch.randn(4096, d)
    torch.set_num_threads(4)
    t0 = time.perf_counter()
    for _ in range(5):
        with torch.no_grad():
            lin(x)
    per_token_s = (time.perf_counter() - t0) / (5 * 4096)

    return {
        "dance_encode_tokens_once": dance_tokens,
        "dance_selection_reads": dance_select_reads,
        "llm_baseline_reencode_tokens": llm_tokens,
        "full_graph_llm_tokens": full_llm_tokens,
        "dance_vs_llm_token_ratio": round(llm_tokens / dance_tokens, 3),
        "dance_vs_fullgraph_ratio": round(full_llm_tokens / dance_tokens, 3),
        "per_token_forward_s": per_token_s,
        "est_dance_encode_time_s": round(dance_tokens * per_token_s, 3),
        "est_llm_baseline_time_s": round(llm_tokens * per_token_s, 3),
        "est_fullgraph_llm_time_s": round(full_llm_tokens * per_token_s, 3),
    }


def main():
    out = {"A_edges_O_Kq": partA_edges(),
           "B_time_scaling": partB_time(),
           "C_encoding_cost": partC_encoding()}
    os.makedirs("outputs", exist_ok=True)
    with open("outputs/claim5_complexity.json", "w") as f:
        json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()