"""MVP verification for Spiderweb Hierarchical Attention.
Run: python3 spiderweb_experiment.py
"""
import math, json, random
import numpy as np


def hyperbolic_distance(dx, y, C=1.0):
    # Half-space metric with equal heights: C arcosh(1+dx^2/(2 y^2)).
    z = 1.0 + (dx * dx) / (2.0 * y * y)
    return C * np.arccosh(np.maximum(z, 1.0))


def preferred_level(dx, n, L, eta, y0, tau):
    levels = np.arange(L + 1)
    ys = y0 * eta ** levels
    ds = hyperbolic_distance(dx, ys)
    return int(np.argmin(np.abs(ds - tau))), ds


def predicted_transition(dx, L, eta, y0, tau):
    # Continuous solution of d(dx,y)=tau, then nearest/clipped dyadic level.
    ystar = dx / math.sqrt(2.0 * (math.cosh(tau) - 1.0))
    raw = math.log(ystar / y0, eta)
    return min(L, max(0, int(round(raw)))), raw


def hierarchy_work(n, neighbor_radius=1):
    # One summary and one fixed-radius horizontal attention neighborhood per level;
    # token-level gates/broadcasts are counted, as required by the pseudocode.
    cells = n
    horizontal = 0
    broadcasts = 0
    levels = 0
    while cells >= 1:
        horizontal += cells * (min(cells - 1, neighbor_radius) * 2 + 1)
        broadcasts += n
        levels += 1
        if cells == 1: break
        cells = (cells + 1) // 2
    return horizontal + broadcasts, levels, horizontal, broadcasts


def graph_coverage(n, kind, window=2, radius=1):
    """Reachable nodes after one communication layer; boolean adjacency."""
    A = np.zeros((n, n), dtype=bool)
    if kind == 'dense':
        A[:] = True
    elif kind == 'local':
        for i in range(n):
            A[i, max(0, i-window):min(n, i+window+1)] = True
    elif kind == 'spiderweb':
        # Undirected approximation to pooling, horizontal neighboring cells,
        # and broadcast through every dyadic level.
        L = int(math.ceil(math.log2(n)))
        for i in range(n):
            for l in range(L + 1):
                size = 2 ** l
                lo = (i // size) * size; hi = min(n, lo + size)
                A[i, lo:hi] = True
                for off in range(1, radius + 1):
                    for c in (lo - off*size, hi + (off-1)*size):
                        if 0 <= c < n:
                            A[i, c:min(n, c+size)] = True
    return A


def main():
    np.random.seed(7); random.seed(7)
    eta, y0, tau, n, L = 2.0, 0.02, 1.0, 256, 8
    # Prediction 1: d(dx,y_l) strictly decreases with level for every dx>0.
    dxs = np.array([1/n, 8/n, 32/n, 100/n])
    monotone = []
    for dx in dxs:
        _, ds = preferred_level(dx, n, L, eta, y0, tau)
        monotone.append(bool(np.all(np.diff(ds) < 0)))
    # Prediction 2: scale transition agrees with closed form inverse metric.
    rows = []
    for dx in np.linspace(1/n, 0.95, 25):
        obs, ds = preferred_level(float(dx), n, L, eta, y0, tau)
        pred, raw = predicted_transition(float(dx), L, eta, y0, tau)
        rows.append((float(dx), obs, pred, abs(obs-pred)))
    errors = np.array([r[3] for r in rows])
    # Prediction 3: work/(n log2 n) stays bounded while dense/(n log n) grows.
    ns = np.array([32,64,128,256,512,1024,2048])
    sw = np.array([hierarchy_work(int(x))[0] for x in ns])
    dense = ns.astype(float)**2
    norm_sw = sw/(ns*np.log2(ns)); norm_dense = dense/(ns*np.log2(ns))
    # Toy coverage and operations at n=256.
    cov = {}
    for kind in ('dense','local','spiderweb'):
        A = graph_coverage(n, kind)
        cov[kind] = {'reachable_from_middle': int(A[n//2].sum()),
                     'mean_reachable': float(A.sum()/n)}
    result = {
      'settings': {'n':n,'eta':eta,'y0':y0,'tau':tau,'levels':L},
      'prediction_1_monotone_distance': {'dx':dxs.tolist(),'confirmed_each':monotone,
          'claim':'equal-height hyperbolic distance strictly decreases as y_l grows'},
      'prediction_2_transition': {'formula':'y*=dx/sqrt(2(cosh(tau)-1)); l*=round(log_eta(y*/y0))',
          'mean_abs_level_error':float(errors.mean()), 'max_abs_level_error':float(errors.max()),
          'fraction_exact':float(np.mean(errors==0)), 'samples':rows[::6]},
      'prediction_3_work_scaling': {'n':ns.tolist(),'work':sw.tolist(),
          'work_over_n_log2n':norm_sw.tolist(),'dense_over_n_log2n':norm_dense.tolist(),
          'spiderweb_ratio_range':[float(norm_sw.min()),float(norm_sw.max())],
          'dense_ratio_growth':float(norm_dense[-1]/norm_dense[0])},
      'toy_communication': {'spiderweb_work':hierarchy_work(n)[0],
          'dense_attention_scores':n*n,'local_attention_scores':n*5,'coverage':cov}
    }
    print(json.dumps(result, indent=2))

if __name__ == '__main__': main()

class SpiderwebAttention:
    """Small readable NumPy implementation of the proposed hierarchy.

    This is an inference/MVP layer: cell pooling is mean pooling, horizontal
    attention uses nearby cells, and each level broadcasts a gated summary.
    """
    def __init__(self, dim, eta=2, levels=None, radius=1, seed=0):
        self.dim, self.eta, self.radius = dim, int(eta), int(radius)
        self.rng = np.random.default_rng(seed)
        self.levels = levels
        self.gates = None

    def __call__(self, x):
        x = np.asarray(x, dtype=float)
        n, d = x.shape
        L = self.levels if self.levels is not None else int(math.ceil(math.log2(n)))
        cur = x.copy()
        outputs = []
        for lev in range(L + 1):
            size = self.eta ** lev
            cells = [(a, min(a + size, n)) for a in range(0, n, size)]
            s = np.stack([cur[a:b].mean(0) for a, b in cells])
            q = s / math.sqrt(d)
            scores = q @ s.T
            mask = np.full(scores.shape, -np.inf)
            for c in range(len(cells)):
                lo, hi = max(0, c-self.radius), min(len(cells), c+self.radius+1)
                mask[c, lo:hi] = scores[c, lo:hi]
            mask -= np.max(mask, axis=1, keepdims=True)
            att = np.exp(mask); att /= att.sum(axis=1, keepdims=True)
            updated = att @ s
            broadcast = np.zeros_like(cur)
            for c, (a, b) in enumerate(cells): broadcast[a:b] = updated[c]
            gate = 1.0 / (lev + 1.0)
            cur = cur + gate * broadcast
            outputs.append(cur.copy())
        return cur, outputs


def dense_mean(x):
    return x + np.broadcast_to(x.mean(0), x.shape)


def local_mean(x, window=2):
    out = x.copy()
    for i in range(len(x)):
        out[i] += x[max(0, i-window):min(len(x), i+window+1)].mean(0)
    return out