import json, math
from pathlib import Path
import numpy as np

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

# Histogram mechanism from the proposal.
def histogram(omega, Amax, R):
    omega = np.asarray(omega, dtype=int)
    q = np.minimum(R - 1, np.floor(R * np.clip(omega, 0, Amax) / Amax).astype(int))
    return np.bincount(q, minlength=R).astype(float) / len(omega)

def quantizer_sweep():
    # Prediction: bin r starts at the smallest integer ceil(r*Amax/R),
    # except bin zero; clipping makes the final bin absorb the endpoint.
    A, R = 11, 4
    observed = []
    for r in range(1, R):
        vals = np.arange(A + 1)
        qs = np.minimum(R-1, np.floor(R*vals/A).astype(int))
        observed.append(int(vals[qs == r][0]))
    predicted = [int(math.ceil(r*A/R)) for r in range(1, R)]
    return {"Amax": A, "R": R, "predicted_first_integer": predicted,
            "observed_first_integer": observed,
            "max_threshold_error": int(max(abs(a-b) for a,b in zip(predicted, observed)))}

def invariance_sweep():
    # Prediction: any permutation of nodes, and any replicated cluster with
    # identical empirical distribution, has exactly the same fixed-size state.
    A, R = 8, 5
    omega = np.array([0, 1, 2, 2, 5, 8, 3, 3, 1, 0])
    h = histogram(omega, A, R)
    perm_errors, replicate_errors = [], []
    for _ in range(100):
        perm_errors.append(float(np.max(np.abs(h - histogram(rng.permutation(omega), A, R)))))
        k = int(rng.integers(1, 8))
        replicate_errors.append(float(np.max(np.abs(h - histogram(np.tile(omega, k), A, R)))))
    # Deliberately show cardinality is not represented: same distribution, sizes differ.
    return {"permutation_max_abs_error": max(perm_errors),
            "replication_max_abs_error": max(replicate_errors),
            "state_dimension": R,
            "sizes_tested": [len(omega), len(omega)*7]}

def synchronous_order_sweep():
    # Two-node linear update x_i <- (1-eta)x_i + eta*x_neighbor.
    # For x=[1,0], synchronous vs sequential (0 then 1) differs exactly eta^2
    # in the second component, demonstrating why pre-update gathering matters.
    etas = np.linspace(.05, .95, 10)
    rows = []
    for eta in etas:
        x = np.array([1., 0.])
        sync = np.array([(1-eta)*x[0] + eta*x[1], (1-eta)*x[1] + eta*x[0]])
        seq = x.copy()
        seq[0] = (1-eta)*seq[0] + eta*seq[1]
        seq[1] = (1-eta)*seq[1] + eta*seq[0]
        err = float(np.max(np.abs(sync-seq)))
        rows.append((float(eta), err, float(eta**2)))
    rel = max(abs(e-p) for _,e,p in rows)
    return {"rows": rows, "predicted_linf_error": "eta^2",
            "max_absolute_error_vs_prediction": rel}

def make_problem(n=96, clusters=8, seed=SEED):
    r = np.random.default_rng(seed)
    # Ring plus random edges; target is a smooth binary-ish signal.
    edges = set()
    for i in range(n):
        for d in (1, 2):
            j=(i+d)%n; edges.add((min(i,j),max(i,j)))
    for _ in range(2*n):
        i,j=r.integers(0,n,2)
        if i != j: edges.add((min(i,j),max(i,j)))
    adj=[set() for _ in range(n)]
    for i,j in edges: adj[i].add(j); adj[j].add(i)
    target=np.sin(np.arange(n)*2*np.pi/24)
    x=r.normal(0,.9,n)
    blocks=[np.arange(a*n//clusters,(a+1)*n//clusters) for a in range(clusters)]
    return adj,target,x,blocks

def residual_mass(x, target, adj):
    # Residual indicator is thresholded local error, as in the proposal.
    return float(np.sum(np.abs(x-target) > .15))

def update_cluster(x, target, adj, ids, eta=.55):
    # All neighbor reads are from x, then scatter: synchronous within cluster.
    old=x.copy(); out=[]
    for i in ids:
        neigh=list(adj[i])
        mean=np.mean(old[neigh]) if neigh else old[i]
        out.append((1-eta)*old[i] + eta*(.55*mean+.45*target[i]))
    x[np.asarray(ids)] = out

def mini_compare():
    adj,target,x0,blocks=make_problem()
    # Standard: one full synchronous sweep is the natural reference.
    full=[]; x=x0.copy()
    for t in range(30):
        old=x.copy()
        for ids in blocks:
            # full sweep uses a single pre-update state for every node
            update_cluster(x, target, adj, np.arange(len(x)), eta=.55)
            break
        full.append(residual_mass(x,target,adj))
    # Idea: observe per-cluster residual histograms and greedily select the
    # cluster with greatest high-bin mass; this is the exploitation behavior
    # a learned Q scheduler is intended to approximate.
    idea=[]; x=x0.copy(); decisions=0
    for t in range(30*len(blocks)):
        scores=[]
        for ids in blocks:
            err=np.abs(x[ids]-target[ids]); omega=np.minimum(8,(err>.15).astype(int)*8)
            scores.append(histogram(omega,8,4)[-1])
        a=int(np.argmax(scores)); update_cluster(x,target,adj,blocks[a],eta=.55); decisions+=1
        if decisions % len(blocks)==0: idea.append(residual_mass(x,target,adj))
    return {"full_synchronous_residual_each_sweep": full,
            "histogram_selected_residual_each_8_decisions": idea,
            "final_full": full[-1], "final_idea": idea[-1],
            "idea_decisions_for_30_sweep_equivalents": decisions,
            "note":"This is a tiny signal check, not a claim of wall-clock speedup."}

def main():
    out={"seed":SEED,"predictions":{
      "P1_histogram_permutation_and_cardinality_invariance": invariance_sweep(),
      "P2_quantizer_transition_thresholds": quantizer_sweep(),
      "P3_synchronous_cluster_update_order_error": synchronous_order_sweep()},
      "mini_experiment":mini_compare()}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))
if __name__ == '__main__': main()